Changed 'upstream_data_directories' field in MachineConfig into a dictionary, and updated functions that use it to reflect this; changed 'find_upstream_visits' to return a nested dictionary where visits are grouped by instrument name

Author: tieneupin

src/murfey/client/tui/screens.py

@@ -685,15 +685,23 @@ def on_button_pressed(self, event: Button.Pressed):
                 self.app.push_screen("launcher")
 
         if machine_data.get("upstream_data_directories"):
-            upstream_downloads = capture_get(
+            upstream_downloads: dict[str, dict[str, Path]] = capture_get(
                 base_url=str(self.app._environment.url.geturl()),
                 router_name="session_control.correlative_router",
                 function_name="find_upstream_visits",
                 token=token,
                 session_id=self.app._environment.murfey_session,
             ).json()
+            # Pass flattened dict for backwards compatibility
             self.app.install_screen(
-                UpstreamDownloads(upstream_downloads), "upstream-downloads"
+                UpstreamDownloads(
+                    {
+                        visit_name: visit_dir
+                        for _, upstream_visits in upstream_downloads.items()
+                        for visit_name, visit_dir in upstream_visits.items()
+                    }
+                ),
+                "upstream-downloads",
             )
             self.app.push_screen("upstream-downloads")
 
@@ -759,15 +767,23 @@ def on_button_pressed(self, event: Button.Pressed):
             self.app.push_screen("directory-select")
 
         if machine_data.get("upstream_data_directories"):
-            upstream_downloads = capture_get(
+            upstream_downloads: dict[str, dict[str, Path]] = capture_get(
                 base_url=str(self.app._environment.url.geturl()),
                 router_name="session_control.correlative_router",
                 function_name="find_upstream_visits",
                 token=token,
                 session_id=self.app._environment.murfey_session,
             ).json()
+            # Pass a flattened dict for backwards compatibility
             self.app.install_screen(
-                UpstreamDownloads(upstream_downloads), "upstream-downloads"
+                UpstreamDownloads(
+                    {
+                        visit_name: visit_dir
+                        for _, upstream_visits in upstream_downloads.items()
+                        for visit_name, visit_dir in upstream_visits.items()
+                    }
+                ),
+                "upstream-downloads",
             )
             self.app.push_screen("upstream-downloads")
 

src/murfey/server/api/shared.py

@@ -147,6 +147,10 @@ def get_foil_hole(session_id: int, fh_name: int, db) -> Dict[str, int]:
 
 
 def find_upstream_visits(session_id: int, db: SQLModelSession):
+    """
+    Returns a nested dictionary, in which visits and the full paths to their directories
+    are further grouped by instrument name.
+    """
     murfey_session = db.exec(
         select(MurfeySession).where(MurfeySession.id == session_id)
     ).one()
@@ -155,12 +159,19 @@ def find_upstream_visits(session_id: int, db: SQLModelSession):
     machine_config = get_machine_config(instrument_name=instrument_name)[
         instrument_name
     ]
-    upstream_visits = {}
+    upstream_visits: dict[str, dict[str, Path]] = {}
     # Iterates through provided upstream directories
-    for p in machine_config.upstream_data_directories:
+    for (
+        upstream_instrument,
+        upstream_data_dir,
+    ) in machine_config.upstream_data_directories.items():
         # Looks for visit name in file path
-        for v in Path(p).glob(f"{visit_name.split('-')[0]}-*"):
-            upstream_visits[v.name] = v / machine_config.processed_directory_name
+        current_upstream_visits = {}
+        for visit_path in Path(upstream_data_dir).glob(f"{visit_name.split('-')[0]}-*"):
+            current_upstream_visits[visit_path.name] = (
+                visit_path / machine_config.processed_directory_name
+            )
+        upstream_visits[upstream_instrument] = current_upstream_visits
     return upstream_visits
 
 
@@ -170,7 +181,7 @@ def get_upstream_tiff_dirs(visit_name: str, instrument_name: str) -> List[Path]:
         instrument_name
     ]
     for directory_name in machine_config.upstream_data_tiff_locations:
-        for p in machine_config.upstream_data_directories:
+        for _, p in machine_config.upstream_data_directories.items():
             if (Path(p) / secure_filename(visit_name)).is_dir():
                 processed_dir = Path(p) / secure_filename(visit_name) / directory_name
                 tiff_dirs.append(processed_dir)

src/murfey/util/config.py

@@ -58,7 +58,7 @@ class MachineConfig(BaseModel):  # type: ignore
     allow_removal: bool = False
 
     # Upstream data download setup
-    upstream_data_directories: list[Path] = []  # Previous sessions
+    upstream_data_directories: dict[str, Path] = {}  # Previous sessions
     upstream_data_download_directory: Optional[Path] = None  # Set by microscope config
     upstream_data_tiff_locations: list[str] = ["processed"]  # Location of CLEM TIFFs