Manage size of zip data

openflexure_microscope/plugins/default/zip_builder/plugin.py

@@ -5,6 +5,7 @@ from openflexure_microscope.devel import (
     request,
     jsonify,
     taskify,
+    update_task_progress,
 )
 
 from flask import send_file, abort
@@ -29,7 +30,7 @@ class ZipBuilderAPIView(MicroscopeViewPlugin):
 
 class ZipListAPIView(MicroscopeViewPlugin):
     def get(self):
-        return jsonify(list(self.plugin.session_zips.keys()))
+        return jsonify(self.plugin.session_zips)
 
 
 class ZipGetterAPIView(MicroscopeViewPlugin):
@@ -40,7 +41,7 @@ class ZipGetterAPIView(MicroscopeViewPlugin):
         logging.info(f"Session ID: {session_id}")
 
         return send_file(
-            self.plugin.session_zips[session_id].name,
+            self.plugin.zip_from_id(session_id).name,
             mimetype="application/zip",
             as_attachment=True,
             attachment_filename=f"{session_id}.zip",
@@ -51,9 +52,9 @@ class ZipGetterAPIView(MicroscopeViewPlugin):
             return abort(404)  # 404 Not Found
 
         logging.info(f"Session ID: {session_id}")
-        logging.debug(self.plugin.session_zips[session_id].name)
 
-        fp = self.plugin.session_zips[session_id]
+        fp = self.plugin.zip_from_id(session_id)
+        logging.debug(fp.name)
         fp.close()
         os.unlink(fp.name)
 
@@ -82,18 +83,53 @@ class ZipBuilderPlugin(MicroscopePlugin):
     def build_zip_from_capture_ids(self, capture_id_list):
         logging.debug(capture_id_list)
 
+        # Get array of captures from IDs
+        capture_list = [
+            self.microscope.camera.image_from_id(capture_id)
+            for capture_id in capture_id_list
+        ]
+        # Remove Nones from list (missing/invalid captures)
+        capture_list = [capture for capture in capture_list if capture]
+
+        # Get size (in bytes) of each capture
+        capture_sizes = [
+            os.path.getsize(capture_obj.file) for capture_obj in capture_list
+        ]
+        # Calculate size of input data in megabytes
+        data_size_megabytes = sum(capture_sizes) * 1e-6
+
+        # If more than 1GB
+        if data_size_megabytes > 1000:
+            # Throw exception
+            raise Exception(
+                "Zip data cannot exceed 1GB. Please transfer data manually."
+            )
+
+        # Number of files to add (used for task progress)
+        n_files = len(capture_id_list)
+
+        # Create temporary file
         fp = tempfile.NamedTemporaryFile(delete=False)
 
+        # Open temp file as a ZIP file
         with zipfile.ZipFile(fp, "w") as zipObj:
-            for capture_id in capture_id_list:
-                capture_obj = self.microscope.camera.image_from_id(capture_id)
+            for index, capture_obj in enumerate(capture_list):
+                # Add to ZIP file if it exists
                 filePath = capture_obj.file
                 zipObj.write(filePath)
+                # Update task progress
+                update_task_progress(int((index / n_files) * 100))
 
         session_id = uuid.uuid4().hex
-        self.session_zips[session_id] = fp
+        # self.session_zips[session_id] = fp
+        self.session_zips[session_id] = {
+            "id": session_id,
+            "fp": fp,
+            "data_size": data_size_megabytes,
+            "zip_size": os.path.getsize(fp.name) * 1e-6,
+        }
 
-        return session_id
+        return self.session_zips[session_id]
 
     def zip_from_id(self, session_id):
-        return self.session_zips[session_id]
+        return self.session_zips[session_id]["fp"]