Generate dictionary of available storage locations

from labthings.server.extensions import BaseExtension

from openflexure_microscope.paths import settings_file_path

from openflexure_microscope.paths import settings_file_path, check_rw

from openflexure_microscope.config import OpenflexureSettingsFile

from openflexure_microscope.camera.base import BASE_CAPTURE_PATH

import logging

import os

import psutil

from sys import platform

AS_SETTINGS_PATH = settings_file_path("autostorage_settings.json")

def get_partitions():

    return [disk.mountpoint for disk in psutil.disk_partitions() if "rw" in disk.opts]

def get_permissive_partitions():

    return [partition for partition in get_partitions() if check_rw(partition)]

def get_permissive_locations():

    return [

        (partition, os.path.join(partition, "openflexure", "data", "micrographs"))

        for partition in get_permissive_partitions()

    ]

def get_current_location(camera):

    return camera.paths.get("default")

def get_default_location():

    return BASE_CAPTURE_PATH

def get_all_locations(camera):

    locations = {"System": get_current_location(camera)}

    # If default is not already listed (e.g. if it's currently set)

    if get_default_location() not in locations.values():

        locations["Default"] = get_default_location()

    for ppartition, plocation in get_permissive_locations():

        pdrive = os.path.splitdrive(plocation)[0]

        if not (

            pdrive  # If path actually has a drive (basically just Windows?)

            and any(  # And shares a common drive with an existing location

                [

                    pdrive == os.path.splitdrive(location)[0]

                    for location in locations.values()

                ]

            )

        ):

            locations[ppartition] = plocation

    # Strip out Nones

    return {k: v for k, v in locations.items() if v}

class CaptureStorageLocation:

    def __init__(self, mountpoint):

        pass

class AutostorageExtension(BaseExtension):

    def __init__(self):

        BaseExtension.__init__(

            self.camera = microscope_obj.camera

            logging.debug(f"Autostorage extension bound to camera {self.camera}")

            logging.debug(self.get_locations())

    def get_locations(self):

        if self.camera:

            return get_all_locations(self.camera)

        else:

            return []

autostorage_extension_v2 = AutostorageExtension()