155 lines
4.4 KiB
Python
155 lines
4.4 KiB
Python
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_component
|
|
from openflexure_microscope.common.flask_labthings.view import View
|
|
from openflexure_microscope.common.flask_labthings.extensions import BaseExtension
|
|
from openflexure_microscope.common.flask_labthings.decorators import (
|
|
ThingAction,
|
|
ThingProperty,
|
|
)
|
|
|
|
|
|
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()
|
|
session_key = str(session_id)
|
|
# self.session_zips[session_id] = fp
|
|
self.session_zips[session_key] = {
|
|
"id": session_id,
|
|
"fp": fp,
|
|
"data_size": data_size_megabytes,
|
|
"zip_size": os.path.getsize(fp.name) * 1e-6,
|
|
}
|
|
|
|
return self.session_zips[session_key]
|
|
|
|
def zip_from_id(self, session_id):
|
|
return self.session_zips[session_id]["fp"]
|
|
|
|
|
|
# Create a global ZIP manager
|
|
default_zip_manager = ZipManager()
|
|
|
|
|
|
@ThingAction
|
|
class ZipBuilderAPIView(View):
|
|
def post(self):
|
|
|
|
ids = list(JsonResponse(request).json)
|
|
microscope = find_component("org.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
|
|
|
|
|
|
@ThingProperty
|
|
class ZipListAPIView(View):
|
|
def get(self):
|
|
return jsonify(default_zip_manager.session_zips)
|
|
|
|
|
|
class ZipGetterAPIView(View):
|
|
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_extension_v2 = BaseExtension("org.openflexure.zipbuilder", version="2.0.0-beta.1")
|
|
|
|
zip_extension_v2.add_view(ZipGetterAPIView, "/get/<string:session_id>")
|
|
zip_extension_v2.add_view(ZipListAPIView, "/get")
|
|
|
|
zip_extension_v2.add_view(ZipBuilderAPIView, "/build")
|