Added draft v2 zip builder plugin

openflexure_microscope/api/default_plugins/zip_builder.py

@@ -0,0 +1,148 @@
+from openflexure_microscope.devel import (
+    JsonResponse,
+    request,
+    jsonify,
+    taskify,
+    update_task_progress,
+)
+
+from flask import send_file, abort
+
+import uuid
+import os
+import zipfile
+import tempfile
+import logging
+
+from openflexure_microscope.common.flask_labthings.find import find_device
+from openflexure_microscope.common.flask_labthings.resource import Resource
+from openflexure_microscope.common.flask_labthings.plugins import BasePlugin
+
+
+class ZipManager:
+    """
+    ZIP-builder manager
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self.session_zips = {}
+
+    def build_zip_from_capture_ids(self, microscope, capture_id_list):
+        logging.debug(capture_id_list)
+
+        # Get array of captures from IDs
+        capture_list = [
+            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 index, capture_obj in enumerate(capture_list):
+                # Add to ZIP file if it exists
+                file_path = capture_obj.file
+                rel_path = os.path.relpath(
+                    file_path, microscope.camera.paths["default"]
+                )
+                zipObj.write(file_path, arcname=rel_path)
+                # Update task progress
+                update_task_progress(int((index / n_files) * 100))
+
+        session_id = uuid.uuid4()
+        # 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 self.session_zips[session_id]
+
+    def zip_from_id(self, session_id):
+        return self.session_zips[session_id]["fp"]
+
+# Create a global ZIP manager
+default_zip_manager = ZipManager()
+
+class ZipBuilderAPIView(Resource):
+    def post(self):
+
+        ids = list(JsonResponse(request).json)
+        microscope = find_device("openflexure_microscope")
+
+        task = taskify(default_zip_manager.build_zip_from_capture_ids)(microscope, ids)
+
+        # Return a handle on the autofocus task
+        return jsonify(task.state), 201
+
+
+class ZipListAPIView(Resource):
+    def get(self):
+        return jsonify(default_zip_manager.session_zips)
+
+
+class ZipGetterAPIView(Resource):
+    def get(self, session_id):
+        if not session_id in default_zip_manager.session_zips:
+            return abort(404)  # 404 Not Found
+
+        logging.info(f"Session ID: {session_id}")
+
+        return send_file(
+            default_zip_manager.zip_from_id(session_id).name,
+            mimetype="application/zip",
+            as_attachment=True,
+            attachment_filename=f"{session_id}.zip",
+        )
+
+    def delete(self, session_id):
+        if not session_id in default_zip_manager.session_zips:
+            return abort(404)  # 404 Not Found
+
+        logging.info(f"Session ID: {session_id}")
+
+        fp = default_zip_manager.zip_from_id(session_id)
+        logging.debug(fp.name)
+        fp.close()
+        os.unlink(fp.name)
+
+        assert not os.path.exists(fp.name)
+
+        del default_zip_manager.session_zips[session_id]
+
+        return jsonify({"return": session_id})
+
+
+zip_plugin_v2 = BasePlugin("zip_builder")
+
+zip_plugin_v2.add_view(ZipGetterAPIView, "/get/<string:session_id>")
+zip_plugin_v2.add_view(ZipListAPIView, "/get")
+
+zip_plugin_v2.add_view(ZipBuilderAPIView, "/build")
+zip_plugin_v2.register_action(ZipBuilderAPIView)
+