diff --git a/openflexure_microscope/plugins/default/zip_builder/plugin.py b/openflexure_microscope/plugins/default/zip_builder/plugin.py index d95a59d5..fdf29851 100644 --- a/openflexure_microscope/plugins/default/zip_builder/plugin.py +++ b/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"]