Example of shared beamline structure - i05

src/dodal/beamlines/i05.py

@@ -1,49 +1,24 @@
-from dodal.beamline_specific_utils.i05_shared import pgm as i05_pgm
-from dodal.common.beamlines.beamline_utils import (
-    device_factory,
-)
+from dodal.beamlines.i05_shared import PREFIX
+from dodal.beamlines.i05_shared import devices as i05_shared_devices
 from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
-from dodal.devices.apple2_undulator import (
-    Apple2,
-    UndulatorGap,
-    UndulatorLockedPhaseAxes,
-)
-from dodal.devices.pgm import PlaneGratingMonochromator
-from dodal.devices.synchrotron import Synchrotron
+from dodal.device_manager import DeviceManager
+from dodal.devices.motors import XYZStage
 from dodal.log import set_beamline as set_log_beamline
-from dodal.utils import BeamlinePrefix, get_beamline_name
+from dodal.utils import get_beamline_name
+
+devices = DeviceManager()
+devices.combine(i05_shared_devices)
 
 BL = get_beamline_name("i05")
-PREFIX = BeamlinePrefix(BL)
 set_log_beamline(BL)
 set_utils_beamline(BL)
 
 
-@device_factory()
-def synchrotron() -> Synchrotron:
-    return Synchrotron()
-
-
-@device_factory()
-def pgm() -> PlaneGratingMonochromator:
-    return i05_pgm()
-
-
-@device_factory()
-def id_gap() -> UndulatorGap:
-    return UndulatorGap(prefix=f"{PREFIX.insertion_prefix}-MO-SERVC-01:")
-
-
-@device_factory()
-def id_phase() -> UndulatorLockedPhaseAxes:
-    return UndulatorLockedPhaseAxes(
-        prefix=f"{PREFIX.insertion_prefix}-MO-SERVC-01:",
-        top_outer="PL",
-        btm_inner="PU",
+@devices.factory()
+def sm() -> XYZStage:
+    return XYZStage(
+        f"{PREFIX.beamline_prefix}-EA-SM-01:",
+        x_infix="SAX",
+        y_infix="SAY",
+        z_infix="SAZ",
     )
-
-
-@device_factory()
-def id() -> Apple2:
-    """i05 insertion device."""
-    return Apple2(id_gap=id_gap(), id_phase=id_phase())

src/dodal/beamlines/i05_1.py

@@ -1,22 +1,24 @@
-from dodal.beamline_specific_utils.i05_shared import pgm as i05_pgm
-from dodal.common.beamlines.beamline_utils import device_factory
+from dodal.beamlines.i05_shared import devices as i05_shared_devices
 from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
-from dodal.devices.pgm import PlaneGratingMonochromator
-from dodal.devices.synchrotron import Synchrotron
+from dodal.device_manager import DeviceManager
+from dodal.devices.motors import XYZStage
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
+devices = DeviceManager()
+devices.combine(i05_shared_devices)
+
 BL = get_beamline_name("i05-1")
-PREFIX = BeamlinePrefix(BL, suffix="J")
+J_PREFIX = BeamlinePrefix(BL, suffix="J")
 set_log_beamline(BL)
 set_utils_beamline(BL)
 
 
-@device_factory()
-def pgm() -> PlaneGratingMonochromator:
-    return i05_pgm()
-
-
-@device_factory()
-def synchrotron() -> Synchrotron:
-    return Synchrotron()
+@devices.factory()
+def sm() -> XYZStage:
+    return XYZStage(
+        f"{J_PREFIX.beamline_prefix}-EA-SM-01:",
+        x_infix="SMX",
+        y_infix="SMY",
+        z_infix="SMZ",
+    )

src/dodal/beamlines/i05_shared.py

@@ -0,0 +1,47 @@
+from dodal.device_manager import DeviceManager
+from dodal.devices.apple2_undulator import (
+    Apple2,
+    UndulatorGap,
+    UndulatorLockedPhaseAxes,
+)
+from dodal.devices.i05.enums import Grating
+from dodal.devices.pgm import PlaneGratingMonochromator
+from dodal.devices.synchrotron import Synchrotron
+from dodal.utils import BeamlinePrefix
+
+devices = DeviceManager()
+
+PREFIX = BeamlinePrefix("i05", "I")
+
+
+@devices.factory()
+def synchrotron() -> Synchrotron:
+    return Synchrotron()
+
+
+@devices.factory()
+def pgm() -> PlaneGratingMonochromator:
+    return PlaneGratingMonochromator(
+        prefix=f"{PREFIX.beamline_prefix}-OP-PGM-01:",
+        grating=Grating,
+    )
+
+
+@devices.factory()
+def id_gap() -> UndulatorGap:
+    return UndulatorGap(prefix=f"{PREFIX.insertion_prefix}-MO-SERVC-01:")
+
+
+@devices.factory()
+def id_phase() -> UndulatorLockedPhaseAxes:
+    return UndulatorLockedPhaseAxes(
+        prefix=f"{PREFIX.insertion_prefix}-MO-SERVC-01:",
+        top_outer="PL",
+        btm_inner="PU",
+    )
+
+
+@devices.factory()
+def id() -> Apple2[UndulatorLockedPhaseAxes]:
+    """i05 insertion device."""
+    return Apple2[UndulatorLockedPhaseAxes](id_gap=id_gap(), id_phase=id_phase())

src/dodal/device_manager.py

@@ -343,6 +343,18 @@ def or_raise(self) -> Self:
         return self
 
 
+def combine_dicts_no_duplicates(dict1, dict2):
+    # Find overlapping keys (keys present in both dictionaries)
+    overlapping_keys = set(dict1) & set(dict2)
+
+    if overlapping_keys:
+        raise ValueError(f"Duplicate keys detected: {', '.join(overlapping_keys)}")
+
+    # If no duplicates, combine the dictionaries
+    combined_dict = {**dict1, **dict2}
+    return combined_dict
+
+
 class DeviceManager:
     """Manager to handle building and connecting interdependent devices"""
 
@@ -355,6 +367,17 @@ def __init__(self):
         self._v1_factories = {}
         self._fixtures = {}
 
+    def combine(self, devices: Self) -> None:
+        self._factories = combine_dicts_no_duplicates(
+            self._factories,
+            devices._factories,  # noqa: SLF001
+        )
+        self._v1_factories = combine_dicts_no_duplicates(
+            self._v1_factories,
+            devices._v1_factories,  # noqa: SLF001
+        )
+        self._fixtures = combine_dicts_no_duplicates(self._fixtures, devices._fixtures)  # noqa: SLF001
+
     def fixture(self, func: Callable[[], T]) -> Callable[[], T]:
         """Add a function that can provide fixtures required by the factories"""
         self._fixtures[func.__name__] = func
@@ -518,10 +541,10 @@ def build_devices(
 
         return DeviceBuildResult(built, errors, connection_specs)
 
-    def __contains__(self, name):
+    def __contains__(self, name) -> bool:
         return name in self._factories or name in self._v1_factories
 
-    def __getitem__(self, name):
+    def __getitem__(self, name) -> DeviceFactory | V1DeviceFactory:
         return self._factories.get(name) or self._v1_factories[name]
 
     def _expand_dependencies(

tests/conftest.py

@@ -1,5 +1,6 @@
 import importlib
 import os
+import sys
 from collections.abc import AsyncGenerator
 from pathlib import Path
 from types import ModuleType
@@ -54,6 +55,12 @@ def module_and_devices_for_beamline(request: pytest.FixtureRequest):
                 factory.cache_clear()
         del bl_mod
 
+        # Remove beamline module from sys.modules to make sure we are testing beamlines
+        # in isolated modules and doesn't leak over to the next one.
+        mods_to_clear = [m for m in sys.modules if m.startswith("dodal.beamlines")]
+        for m in mods_to_clear:
+            sys.modules.pop(m, None)
+
 
 def mock_beamline_module_filepaths(bl_name: str, bl_module: ModuleType):
     if mock_attributes := mock_attributes_table.get(bl_name):

tests/test_utils.py

@@ -1,6 +1,7 @@
 import os
 from collections.abc import Iterable, Mapping
 from shutil import copytree
+from types import ModuleType
 from typing import Any, cast
 from unittest.mock import ANY, MagicMock, Mock, patch
 
@@ -308,15 +309,17 @@ def device_b() -> Motor:
 
 
 @pytest.mark.parametrize("bl", ["", "$%^&*", "nonexistent"])
-def test_invalid_beamline_variable_causes_get_device_module_to_raise(bl):
+def test_invalid_beamline_variable_causes_get_device_module_to_raise(bl: str):
     with patch.dict(os.environ, {"BEAMLINE": bl}), pytest.raises(ValueError):
         get_beamline_based_on_environment_variable()
 
 
 @pytest.mark.parametrize("bl,module", [("i04", i04), ("i23", i23)])
-def test_valid_beamline_variable_causes_get_device_module_to_return_module(bl, module):
+def test_valid_beamline_variable_causes_get_device_module_to_return_module(
+    bl: str, module: ModuleType
+):
     with patch.dict(os.environ, {"BEAMLINE": bl}):
-        assert get_beamline_based_on_environment_variable() == module
+        assert get_beamline_based_on_environment_variable().__name__ == module.__name__
 
 
 def test_find_next_run_number_from_files_gets_correct_number():