Commit 2a8314f3 authored by Joel Collins's avatar Joel Collins
Browse files

Manage size of zip data

parent 98d8993b

openflexure_microscope/plugins/default/zip_builder/plugin.py

+45 −9
Original line number Diff line number Diff line
@@ -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"]