diff --git a/openflexure_microscope/api/default_plugins/zip_builder.py b/openflexure_microscope/api/default_plugins/zip_builder.py new file mode 100644 index 00000000..483d694e --- /dev/null +++ b/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/") +zip_plugin_v2.add_view(ZipListAPIView, "/get") + +zip_plugin_v2.add_view(ZipBuilderAPIView, "/build") +zip_plugin_v2.register_action(ZipBuilderAPIView) +