refactor: Wrap Devices in BlueapiDevice (#1744)

This PR follows #1739. Devices in `BlueskyContext` are now wrapped in a
`BlueapiDevice` dataclass. This is similar to #1743.

Summary of changes:

- `BlueskyContext.devices` was `dict[str, Device]` where Device
represented bluesky protocols. Now `BlueskyContext.devices` is a
`dict[str, BlueapiDevice]`
- `DeviceModel` and the REST side are unchanged
- `ctx.devices` is a dictionary of `BlueapiDevice` wrappers
- `find_device` will look up the first part of the address (`stage` in
`stage.x`) in `ctx.devices` and returns a `Device`
- `find_component` is no longer a standalone function but has been moved
to `BlueapiDevice`. This functions follows from `find_device`, e.g.
finds the `x` on `stage` to give `stage.x`

---------

Co-authored-by: Peter Holloway <peter.holloway@diamond.ac.uk>

Author: shree-iyengar-dls

src/blueapi/core/bluesky_types.py

@@ -131,3 +131,22 @@ def watch(self, __func: Callable) -> None:
         """
 
         ...
+
+
+@dataclass
+class BlueapiDevice:
+    """Wrapper around Device held in the context.
+    Also contains find_component function to look up child devices.
+    At the moment this just holds the Device but will later hold child device trees"""
+
+    device: Device
+
+    def find_component(self, addr: list[str]) -> Device | None:
+        dev = self.device
+        for name in addr:
+            dev = getattr(dev, name, None)
+            if dev is None:
+                return None
+            elif not is_bluesky_compatible_device(dev):
+                raise ValueError(f"Found {dev} but it is not a device")
+        return dev

src/blueapi/core/context.py

@@ -42,13 +42,13 @@
 from .bluesky_types import (
     BLUESKY_PROTOCOLS,
     AsyncDevice,
+    BlueapiDevice,
     Device,
     Plan,
     PlanGenerator,
     is_bluesky_compatible_device,
     is_bluesky_plan_generator,
 )
-from .device_lookup import find_component
 
 LOGGER = logging.getLogger(__name__)
 
@@ -128,7 +128,7 @@ class BlueskyContext:
     numtracker: NumtrackerClient | None = field(default=None, init=False, repr=False)
     path_provider: PathProvider | None = None
     plans: dict[str, Plan] = field(default_factory=dict)
-    devices: dict[str, Device] = field(default_factory=dict)
+    devices: dict[str, BlueapiDevice] = field(default_factory=dict)
 
     _reference_cache: dict[type, type] = field(default_factory=dict)
     _warned_name_clash: bool = field(default=False, init=False, repr=False)
@@ -196,10 +196,12 @@ def find_device(self, addr: str | list[str]) -> Device | None:
         """
 
         if isinstance(addr, str):
-            list_addr = list(addr.split("."))
-            return self.find_device(list_addr)
-        else:
-            return find_component(self.devices, addr)
+            return self.find_device(addr.split("."))
+
+        head, *tail = addr
+        if entry := self.devices.get(head):
+            return entry.find_component(tail)
+        return None
 
     def with_config(self, config: EnvironmentConfig) -> None:
         if config.metadata is not None:
@@ -381,8 +383,8 @@ def _register_children(self, device: Device) -> None:
     def _add_device(self, name: str, device: Device) -> bool:
         """Add a device unless a different one is already registered as name."""
 
-        existing = self.devices.setdefault(name, device)
-        if existing is not device:
+        existing = self.devices.setdefault(name, BlueapiDevice(device))
+        if existing.device is not device:
             # Expected for e.g. a motor's user_readback, which has the motor's name.
             # The caller skips its children too, as their names start with this one
             if not self._warned_name_clash:
@@ -448,11 +450,12 @@ def __get_pydantic_json_schema__(
                     json_schema = handler(core_schema)
                     json_schema = handler.resolve_ref_schema(json_schema)
                     json_schema["type"] = qualified_name(target)
+                    # devices are wrapped in BlueapiDevice, so check entry.device
                     json_schema["enum"] = sorted(
                         [
                             name
-                            for name, device in self.devices.items()
-                            if isinstance(device, cls.origin or target)
+                            for name, entry in self.devices.items()
+                            if isinstance(entry.device, cls.origin or target)
                         ]
                     )
                     if cls.args:

src/blueapi/core/device_lookup.py

@@ -191,7 +191,9 @@ def get_plan(name: str) -> PlanModel:
 
 def get_devices() -> list[DeviceModel]:
     """Get all available devices in the BlueskyContext"""
-    return [DeviceModel.from_device(device) for device in context().devices.values()]
+    return [
+        DeviceModel.from_device(entry.device) for entry in context().devices.values()
+    ]
 
 
 def get_device(name: str) -> DeviceModel:

src/blueapi/service/interface.py

@@ -50,6 +50,7 @@
     TiledConfig,
 )
 from blueapi.core import BlueskyContext, is_bluesky_compatible_device
+from blueapi.core.bluesky_types import BlueapiDevice
 from blueapi.core.context import generic_bounds, qualified_name
 from blueapi.core.protocols import DeviceConnectResult, DeviceManager
 from blueapi.utils.invalid_config_error import InvalidConfigError
@@ -213,8 +214,8 @@ def demo_plan(foo: int, mov: Movable) -> MsgGenerator:
     # Which child signals exist depends on the upstream Motor class
     expected_movables = sorted(
         name
-        for name, device in devicey_context.devices.items()
-        if isinstance(device, Movable)
+        for name, entry in devicey_context.devices.items()
+        if isinstance(entry.device, Movable)
     )
     assert {"sim", "sim-velocity"} <= set(expected_movables)
     assert schema["properties"] == {
@@ -277,14 +278,14 @@ def test_only_plans_from_all_in_module_detected(empty_context: BlueskyContext):
 
 def test_add_named_device(empty_context: BlueskyContext, sim_motor: Motor):
     empty_context.register_device(sim_motor)
-    assert empty_context.devices[SIM_MOTOR_NAME] is sim_motor
+    assert empty_context.devices[SIM_MOTOR_NAME].device is sim_motor
 
 
 def test_add_nameless_device(
     empty_context: BlueskyContext, some_configurable: SomeConfigurable
 ):
     empty_context.register_device(some_configurable, "conf")
-    assert empty_context.devices["conf"] is some_configurable
+    assert empty_context.devices["conf"].device is some_configurable
 
 
 def test_add_nameless_device_without_override(
@@ -297,7 +298,7 @@ def test_add_nameless_device_without_override(
 
 def test_override_device_name(empty_context: BlueskyContext, sim_motor: Motor):
     empty_context.register_device(sim_motor, "foo")
-    assert empty_context.devices["foo"] is sim_motor
+    assert empty_context.devices["foo"].device is sim_motor
 
 
 class Axis(Device):
@@ -336,9 +337,11 @@ def test_register_device_registers_children_recursively(
         "table-mount-y-position",
         "table-mount-y-velocity",
     }
-    assert empty_context.devices["table-mount"] is table.mount
-    assert empty_context.devices["table-mount-x"] is table.mount.x
-    assert empty_context.devices["table-mount-x-velocity"] is table.mount.x.velocity
+    assert empty_context.devices["table-mount"].device is table.mount
+    assert empty_context.devices["table-mount-x"].device is table.mount.x
+    assert (
+        empty_context.devices["table-mount-x-velocity"].device is table.mount.x.velocity
+    )
 
 
 def test_child_sharing_parent_name_is_skipped_and_warned_once(
@@ -350,7 +353,7 @@ def test_child_sharing_parent_name_is_skipped_and_warned_once(
     empty_context.register_device(motor)
     empty_context.register_device(SimMotor(name="other"))
 
-    assert empty_context.devices["motor"] is motor
+    assert empty_context.devices["motor"].device is motor
     assert motor.user_readback not in empty_context.devices.values()
     assert empty_context.find_device("motor.user_readback") is motor.user_readback
     records = caplog.get_records("call")
@@ -373,9 +376,9 @@ def test_device_vector_children_are_registered(empty_context: BlueskyContext):
     bank = AxisBank(name="bank")
     empty_context.register_device(bank)
 
-    assert empty_context.devices["bank-axes"] is bank.axes
-    assert empty_context.devices["bank-axes-1"] is bank.axes[1]
-    assert empty_context.devices["bank-axes-2-velocity"] is bank.axes[2].velocity
+    assert empty_context.devices["bank-axes"].device is bank.axes
+    assert empty_context.devices["bank-axes-1"].device is bank.axes[1]
+    assert empty_context.devices["bank-axes-2-velocity"].device is bank.axes[2].velocity
 
 
 def test_overridden_registration_name_does_not_rename_children(
@@ -386,9 +389,9 @@ def test_overridden_registration_name_does_not_rename_children(
     motor = SimMotor(name="motor")
     empty_context.register_device(motor, name="foo")
 
-    assert empty_context.devices["foo"] is motor
-    assert empty_context.devices["motor"] is motor.user_readback
-    assert empty_context.devices["motor-velocity"] is motor.velocity
+    assert empty_context.devices["foo"].device is motor
+    assert empty_context.devices["motor"].device is motor.user_readback
+    assert empty_context.devices["motor-velocity"].device is motor.velocity
 
 
 def test_unrelated_name_collision_keeps_existing_device_and_logs(
@@ -401,9 +404,9 @@ def test_unrelated_name_collision_keeps_existing_device_and_logs(
     stage = Stage(name="stage")
     empty_context.register_device(stage)
 
-    assert empty_context.devices["stage-x"] is sentinel
+    assert empty_context.devices["stage-x"].device is sentinel
     assert "stage-x-position" not in empty_context.devices
-    assert empty_context.devices["stage-y"] is stage.y
+    assert empty_context.devices["stage-y"].device is stage.y
     assert any("stage-x" in record.message for record in caplog.get_records("call"))
 
 
@@ -417,7 +420,7 @@ def test_top_level_name_collision_keeps_existing_device_and_logs(
     motor = SimMotor(name="stage-x")
     empty_context.register_device(motor)
 
-    assert empty_context.devices["stage-x"] is stage.x
+    assert empty_context.devices["stage-x"].device is stage.x
     assert motor.velocity not in empty_context.devices.values()
     assert any("stage-x" in record.message for record in caplog.get_records("call"))
 
@@ -898,7 +901,7 @@ def test_single_device_manager(empty_context: BlueskyContext):
     sdm = StaticDeviceManager(devices={"foo": foo})
     empty_context.with_device_manager(sdm)
 
-    assert empty_context.devices == {"foo": foo}
+    assert empty_context.devices == {"foo": BlueapiDevice(foo)}
 
 
 def test_device_manager_build_error(empty_context: BlueskyContext):
@@ -968,7 +971,7 @@ def test_device_manager_environment_source(empty_context: BlueskyContext):
     with patch("blueapi.core.context.import_module") as imp_mod:
         imp_mod.side_effect = lambda mod: dev_mod if mod == "foo.bar" else None
         empty_context.with_config(env)
-    assert empty_context.devices == {"foo": foo}
+    assert empty_context.devices == {"foo": BlueapiDevice(foo)}
 
 
 def test_non_device_manager_errors(empty_context: BlueskyContext):