Load Contexts in Analyser via entry points

Dockerfiles/murfey-instrument-server

@@ -32,7 +32,7 @@ RUN apt-get update && \
         pip \
         build \
         importlib-metadata && \
-    /venv/bin/python -m pip install /python-murfey
+    /venv/bin/python -m pip install /python-murfey[sxt]
 
 
 # Transfer completed Murfey build to base image

pyproject.toml

@@ -96,6 +96,15 @@ GitHub = "https://github.com/DiamondLightSource/python-murfey"
 "murfey.transfer" = "murfey.cli.transfer:run"
 [project.entry-points."murfey.config.extraction"]
 "murfey_machine" = "murfey.util.config:get_extended_machine_config"
+[project.entry-points."murfey.contexts"]
+AtlasContext = "murfey.client.contexts.atlas:AtlasContext"
+CLEMContext = "murfey.client.contexts.clem:CLEMContext"
+FIBContext = "murfey.client.contexts.fib:FIBContext"
+SPAContext = "murfey.client.contexts.spa:SPAContext"
+SPAMetadataContext = "murfey.client.contexts.spa_metadata:SPAMetadataContext"
+SXTContext = "murfey.client.contexts.sxt:SXTContext"
+TomographyContext = "murfey.client.contexts.tomo:TomographyContext"
+TomographyMetadataContext = "murfey.client.contexts.tomo_metadata:TomographyMetadataContext"
 [project.entry-points."murfey.workflows"]
 "atlas_update" = "murfey.workflows.register_atlas_update:run"
 "clem.align_and_merge" = "murfey.workflows.clem.align_and_merge:run"

src/murfey/client/analyser.py

@@ -8,21 +8,15 @@
 
 from __future__ import annotations
 
+import functools
 import logging
 import queue
 import threading
+from importlib.metadata import entry_points
 from pathlib import Path
 from typing import Type
 
 from murfey.client.context import Context
-from murfey.client.contexts.atlas import AtlasContext
-from murfey.client.contexts.clem import CLEMContext
-from murfey.client.contexts.fib import FIBContext
-from murfey.client.contexts.spa import SPAModularContext
-from murfey.client.contexts.spa_metadata import SPAMetadataContext
-from murfey.client.contexts.sxt import SXTContext
-from murfey.client.contexts.tomo import TomographyContext
-from murfey.client.contexts.tomo_metadata import TomographyMetadataContext
 from murfey.client.destinations import find_longest_data_directory
 from murfey.client.instance_environment import MurfeyInstanceEnvironment
 from murfey.client.rsync import RSyncerUpdate, TransferResult
@@ -33,6 +27,23 @@
 logger = logging.getLogger("murfey.client.analyser")
 
 
+# Load the Context entry points as a list upon initialisation
+context_eps = list(entry_points(group="murfey.contexts"))
+
+
+@functools.lru_cache(maxsize=1)
+def _get_context(name: str):
+    """
+    Load the desired context from the configured list of entry points.
+    Returns None if the entry point is not found
+    """
+    if context := [ep for ep in context_eps if ep.name == name]:
+        return context[0]
+    else:
+        logger.warning(f"Could not find entry point for {name!r}")
+        return None
+
+
 class Analyser(Observer):
     def __init__(
         self,
@@ -145,7 +156,9 @@ def _find_context(self, file_path: Path) -> bool:
                 )
             )
         ):
-            self._context = CLEMContext(
+            if (context := _get_context("CLEMContext")) is None:
+                return False
+            self._context = context.load()(
                 "leica",
                 self._basepath,
                 self._murfey_config,
@@ -166,7 +179,9 @@ def _find_context(self, file_path: Path) -> bool:
                 and "Sites" in file_path.parts
             )
         ):
-            self._context = FIBContext(
+            if (context := _get_context("FIBContext")) is None:
+                return False
+            self._context = context.load()(
                 "autotem",
                 self._basepath,
                 self._murfey_config,
@@ -183,7 +198,9 @@ def _find_context(self, file_path: Path) -> bool:
                 all(path in file_path.parts for path in ("LayersData", "Layer"))
             )
         ):
-            self._context = FIBContext(
+            if (context := _get_context("FIBContext")) is None:
+                return False
+            self._context = context.load()(
                 "maps",
                 self._basepath,
                 self._murfey_config,
@@ -196,7 +213,9 @@ def _find_context(self, file_path: Path) -> bool:
             # Image metadata stored in "features.json" file
             file_path.name == "features.json" or ()
         ):
-            self._context = FIBContext(
+            if (context := _get_context("FIBContext")) is None:
+                return False
+            self._context = context.load()(
                 "meteor",
                 self._basepath,
                 self._murfey_config,
@@ -208,7 +227,9 @@ def _find_context(self, file_path: Path) -> bool:
         # SXT workflow checks
         # -----------------------------------------------------------------------------
         if file_path.suffix in (".txrm", ".xrm"):
-            self._context = SXTContext(
+            if (context := _get_context("SXTContext")) is None:
+                return False
+            self._context = context.load()(
                 "zeiss",
                 self._basepath,
                 self._murfey_config,
@@ -220,7 +241,9 @@ def _find_context(self, file_path: Path) -> bool:
         # Tomography and SPA workflow checks
         # -----------------------------------------------------------------------------
         if "atlas" in file_path.parts:
-            self._context = AtlasContext(
+            if (context := _get_context("AtlasContext")) is None:
+                return False
+            self._context = context.load()(
                 "serialem" if self._serialem else "epu",
                 self._basepath,
                 self._murfey_config,
@@ -229,7 +252,9 @@ def _find_context(self, file_path: Path) -> bool:
             return True
 
         if "Metadata" in file_path.parts or file_path.name == "EpuSession.dm":
-            self._context = SPAMetadataContext(
+            if (context := _get_context("SPAMetadataContext")) is None:
+                return False
+            self._context = context.load()(
                 "epu",
                 self._basepath,
                 self._murfey_config,
@@ -242,7 +267,9 @@ def _find_context(self, file_path: Path) -> bool:
             or "Thumbnails" in file_path.parts
             or file_path.name == "Session.dm"
         ):
-            self._context = TomographyMetadataContext(
+            if (context := _get_context("TomographyMetadataContext")) is None:
+                return False
+            self._context = context.load()(
                 "tomo",
                 self._basepath,
                 self._murfey_config,
@@ -263,7 +290,9 @@ def _find_context(self, file_path: Path) -> bool:
             ]:
                 if not self._context:
                     logger.info("Acquisition software: EPU")
-                    self._context = SPAModularContext(
+                    if (context := _get_context("SPAContext")) is None:
+                        return False
+                    self._context = context.load()(
                         "epu",
                         self._basepath,
                         self._murfey_config,
@@ -282,7 +311,9 @@ def _find_context(self, file_path: Path) -> bool:
             ):
                 if not self._context:
                     logger.info("Acquisition software: tomo")
-                    self._context = TomographyContext(
+                    if (context := _get_context("TomographyContext")) is None:
+                        return False
+                    self._context = context.load()(
                         "tomo",
                         self._basepath,
                         self._murfey_config,
@@ -322,24 +353,26 @@ def _analyse(self):
                     or transferred_file.name == "EpuSession.dm"
                     and not self._context
                 ):
-                    self._context = SPAMetadataContext(
-                        "epu",
-                        self._basepath,
-                        self._murfey_config,
-                        self._token,
-                    )
+                    if context := _get_context("SPAMetadataContext"):
+                        self._context = context.load()(
+                            "epu",
+                            self._basepath,
+                            self._murfey_config,
+                            self._token,
+                        )
                 elif (
                     "Batch" in transferred_file.parts
                     or "SearchMaps" in transferred_file.parts
                     or transferred_file.name == "Session.dm"
                     and not self._context
                 ):
-                    self._context = TomographyMetadataContext(
-                        "tomo",
-                        self._basepath,
-                        self._murfey_config,
-                        self._token,
-                    )
+                    if context := _get_context("TomographyMetadataContext"):
+                        self._context = context.load()(
+                            "tomo",
+                            self._basepath,
+                            self._murfey_config,
+                            self._token,
+                        )
                 self.post_transfer(transferred_file)
             else:
                 dc_metadata = {}
@@ -364,12 +397,10 @@ def _analyse(self):
                     elif transferred_file.suffix == ".mdoc":
                         mdoc_for_reading = transferred_file
                 if not self._context:
-                    valid_extension = self._find_extension(transferred_file)
-                    if not valid_extension:
+                    if not self._find_extension(transferred_file):
                         logger.error(f"No extension found for {transferred_file}")
                         continue
-                    found = self._find_context(transferred_file)
-                    if not found:
+                    if not self._find_context(transferred_file):
                         logger.debug(
                             f"Couldn't find context for {str(transferred_file)!r}"
                         )
@@ -386,7 +417,7 @@ def _analyse(self):
                             )
                         except Exception as e:
                             logger.error(f"Exception encountered: {e}")
-                        if not isinstance(self._context, AtlasContext):
+                        if "AtlasContext" not in str(self._context):
                             if not dc_metadata:
                                 try:
                                     dc_metadata = self._context.gather_metadata(
@@ -417,31 +448,27 @@ def _analyse(self):
                                 )
                                 self.notify(dc_metadata)
 
-                # If a file with a CLEM context is identified, immediately post it
-                elif isinstance(self._context, CLEMContext):
+                # Contexts that can be immediately posted without additional work
+                elif "CLEMContext" in str(self._context):
                     logger.debug(
-                        f"File {transferred_file.name!r} will be processed as part of CLEM workflow"
+                        f"File {transferred_file.name!r} is part of CLEM workflow"
                     )
                     self.post_transfer(transferred_file)
-
-                elif isinstance(self._context, FIBContext):
+                elif "FIBContext" in str(self._context):
                     logger.debug(
-                        f"File {transferred_file.name!r} will be processed as part of the FIB workflow"
+                        f"File {transferred_file.name!r} is part of the FIB workflow"
                     )
                     self.post_transfer(transferred_file)
-
-                elif isinstance(self._context, SXTContext):
+                elif "SXTContext" in str(self._context):
                     logger.debug(f"File {transferred_file.name!r} is an SXT file")
                     self.post_transfer(transferred_file)
-
-                elif isinstance(self._context, AtlasContext):
+                elif "AtlasContext" in str(self._context):
                     logger.debug(f"File {transferred_file.name!r} is part of the atlas")
                     self.post_transfer(transferred_file)
 
                 # Handle files with tomography and SPA context differently
                 elif not self._extension or self._unseen_xml:
-                    valid_extension = self._find_extension(transferred_file)
-                    if not valid_extension:
+                    if not self._find_extension(transferred_file):
                         logger.error(f"No extension found for {transferred_file}")
                         continue
                     if self._extension:
@@ -480,14 +507,14 @@ def _analyse(self):
                                 self._context._acquisition_software
                             )
                             self.notify(dc_metadata)
-                elif isinstance(
-                    self._context,
-                    (
-                        SPAModularContext,
-                        SPAMetadataContext,
-                        TomographyContext,
-                        TomographyMetadataContext,
-                    ),
+                elif any(
+                    context in str(self._context)
+                    for context in (
+                        "SPAContext",
+                        "SPAMetadataContext",
+                        "TomographyContext",
+                        "TomographyMetadataContext",
+                    )
                 ):
                     context = str(self._context).split(" ")[0].split(".")[-1]
                     logger.debug(

src/murfey/client/contexts/spa.py

@@ -56,7 +56,7 @@ def _get_xml_list_index(key: str, xml_list: list) -> int:
     raise ValueError(f"Key not found in XML list: {key}")
 
 
-class SPAModularContext(Context):
+class SPAContext(Context):
     user_params = [
         ProcessingParameter(
             "dose_per_frame",

src/murfey/client/multigrid_control.py

@@ -11,7 +11,7 @@
 
 from murfey.client.analyser import Analyser
 from murfey.client.context import ensure_dcg_exists
-from murfey.client.contexts.spa import SPAModularContext
+from murfey.client.contexts.spa import SPAContext
 from murfey.client.contexts.tomo import TomographyContext
 from murfey.client.destinations import determine_default_destination
 from murfey.client.instance_environment import MurfeyInstanceEnvironment
@@ -560,7 +560,7 @@ def _start_dc(self, metadata_json):
             )
             log.info("Tomography processing flushed")
 
-        elif isinstance(context, SPAModularContext):
+        elif isinstance(context, SPAContext):
             if self._environment.visit in source.parts:
                 metadata_source = source
             else:

tests/client/test_analyser.py

@@ -6,7 +6,7 @@
 from murfey.client.contexts.atlas import AtlasContext
 from murfey.client.contexts.clem import CLEMContext
 from murfey.client.contexts.fib import FIBContext
-from murfey.client.contexts.spa import SPAModularContext
+from murfey.client.contexts.spa import SPAContext
 from murfey.client.contexts.spa_metadata import SPAMetadataContext
 from murfey.client.contexts.sxt import SXTContext
 from murfey.client.contexts.tomo import TomographyContext
@@ -28,8 +28,8 @@
     ["visit/Batch/BatchPositionsList.xml", TomographyMetadataContext],
     ["visit/Thumbnails/file.mrc", TomographyMetadataContext],
     # SPA
-    ["visit/FoilHole_01234_fractions.tiff", SPAModularContext],
-    ["visit/FoilHole_01234_EER.eer", SPAModularContext],
+    ["visit/FoilHole_01234_fractions.tiff", SPAContext],
+    ["visit/FoilHole_01234_EER.eer", SPAContext],
     # SPA metadata
     ["atlas/atlas.mrc", AtlasContext],
     ["visit/EpuSession.dm", SPAMetadataContext],
@@ -116,7 +116,7 @@ def test_find_context(file_and_context, tmp_path):
     # Checks for the specific workflow contexts
     if isinstance(analyser._context, TomographyContext):
         assert analyser.parameters_model == ProcessingParametersTomo
-    if isinstance(analyser._context, SPAModularContext):
+    if isinstance(analyser._context, SPAContext):
         assert analyser.parameters_model == ProcessingParametersSPA