Refactor get_asset_file_or_stream

OpenDroneMap · May 11, 2024 · 8468fdf · 8468fdf
1 parent 35fc60a
commit 8468fdf
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 10 deletions.
diff --git a/app/api/tasks.py b/app/api/tasks.py
@@ -365,19 +365,20 @@ def get(self, request, pk=None, project_pk=None, asset=""):
 
         # Check and download
         try:
-            asset_fs, is_zipstream = task.get_asset_file_or_zipstream(asset)
+            asset_fs = task.get_asset_file_or_stream(asset)
         except FileNotFoundError:
             raise exceptions.NotFound(_("Asset does not exist"))
 
-        if not is_zipstream and not os.path.isfile(asset_fs):
+        is_stream = not isinstance(asset_fs, str) 
+        if not is_stream and not os.path.isfile(asset_fs):
             raise exceptions.NotFound(_("Asset does not exist"))
 
         download_filename = request.GET.get('filename', get_asset_download_filename(task, asset))
 
-        if not is_zipstream:
-            return download_file_response(request, asset_fs, 'attachment', download_filename=download_filename)
-        else:
+        if is_stream:
             return download_file_stream(request, asset_fs, 'attachment', download_filename=download_filename)
+        else:
+            return download_file_response(request, asset_fs, 'attachment', download_filename=download_filename)
 
 """
 Raw access to the task's asset folder resources

diff --git a/app/models/task.py b/app/models/task.py
@@ -450,16 +450,16 @@ def duplicate(self, set_new_name=True):
 
         return False
 
-    def get_asset_file_or_zipstream(self, asset):
+    def get_asset_file_or_stream(self, asset):
         """
         Get a stream to an asset
         :param asset: one of ASSETS_MAP keys
-        :return: (path|stream, is_zipstream:bool)
+        :return: (path|stream)
         """
         if asset in self.ASSETS_MAP:
             value = self.ASSETS_MAP[asset]
             if isinstance(value, str):
-                return self.assets_path(value), False
+                return self.assets_path(value)
 
             elif isinstance(value, dict):
                 if 'deferred_path' in value and 'deferred_compress_dir' in value:
@@ -469,7 +469,7 @@ def get_asset_file_or_zipstream(self, asset):
                         paths = [p for p in paths if os.path.basename(p['fs']) not in value['deferred_exclude_files']]
                     if len(paths) == 0:
                         raise FileNotFoundError("No files available for download")
-                    return zipfly.ZipStream(paths), True
+                    return zipfly.ZipStream(paths)
                 else:
                     raise FileNotFoundError("{} is not a valid asset (invalid dict values)".format(asset))
             else:

diff --git a/coreplugins/dronedb/api_views.py b/coreplugins/dronedb/api_views.py
@@ -343,7 +343,7 @@ def post(self, request, pk):
 
         settings = get_settings(request)
 
-        available_assets = [task.get_asset_file_or_zipstream(f)[0] for f in list(set(task.available_assets) & set(DRONEDB_ASSETS))]
+        available_assets = [task.get_asset_file_or_stream(f) for f in list(set(task.available_assets) & set(DRONEDB_ASSETS))]
 
         if 'textured_model.zip' in task.available_assets:            
             texture_files = [join(task.assets_path('odm_texturing'), f) for f in listdir(task.assets_path('odm_texturing')) if isfile(join(task.assets_path('odm_texturing'), f))]