refactor: Wrap Devices in BlueapiDevice

src/blueapi/core/bluesky_types.py

@@ -1,5 +1,6 @@
 import inspect
 from collections.abc import Callable, Generator, Mapping
+from dataclasses import dataclass
 from typing import (
     Any,
     Protocol,
@@ -132,3 +133,11 @@ def watch(self, __func: Callable) -> None:
         """
 
         ...
+
+
+@dataclass
+class BlueapiDevice:
+    """Wrapper around Device held in the context.
+    At the moment this just holds the Device but will later hold child device trees"""
+
+    device: Device

src/blueapi/core/context.py

@@ -42,6 +42,7 @@
 from .bluesky_types import (
     BLUESKY_PROTOCOLS,
     AsyncDevice,
+    BlueapiDevice,
     Device,
     Plan,
     PlanGenerator,
@@ -67,20 +68,26 @@ def is_compatible(val: Device, target: type, args: tuple[type, ...] | None):
 
 
 def generic_bounds(val: Device, target: type) -> tuple[type, ...]:
-    for base in getattr(val, "__orig_bases__", ()):
+    unwrapped = val.device if isinstance(val, BlueapiDevice) else val
+    for base in getattr(unwrapped, "__orig_bases__", ()):
         if (get_origin(base) or base) == target:
             return get_args(base)
     return ()
 
 
-def is_compatible_args(val: Device, target: type, args: tuple[type, ...] | None):
+def is_compatible_args(
+    val: Device | BlueapiDevice, target: type, args: tuple[type, ...] | None
+):
+    unwrapped = val.device if isinstance(val, BlueapiDevice) else val
     return (not args) or all(
         actual is Any
         or type(actual) is TypeVar
         or type(expected) is TypeVar
         or expected == actual
         or issubclass(actual, expected)
-        for expected, actual in zip(args, generic_bounds(val, target), strict=False)
+        for expected, actual in zip(
+            args, generic_bounds(unwrapped, target), strict=False
+        )
     )
 
 
@@ -128,7 +135,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)
     plan_functions: dict[str, PlanGenerator] = field(default_factory=dict)
 
     _reference_cache: dict[type, type] = field(default_factory=dict)
@@ -196,10 +203,16 @@ 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("."))
+        # self.devices stores BlueapiDevice wrappers: unwrap the top-level device here
+        # find_component will then look up child devices (like the 'x' in 'motor.x') as
+        # attributes of raw device
+        head, tail = addr[0], addr[1:]
+        entry = self.devices.get(head)
+        device = entry.device if entry else None
+        if tail:
+            return find_component(device, tail)
+        return device
 
     def with_config(self, config: EnvironmentConfig) -> None:
         if config.metadata is not None:
@@ -358,7 +371,7 @@ def register_device(self, device: Device, name: str | None = None) -> None:
             else:
                 raise KeyError(f"Must supply a name for this device: {device}")
 
-        self.devices[name] = device
+        self.devices[name] = BlueapiDevice(device)
 
     def unregister_all_devices(self):
         """Unregister all devices from the context."""
@@ -409,11 +422,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/service/interface.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:

tests/unit_tests/core/test_context.py

@@ -46,6 +46,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
@@ -244,14 +245,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(
@@ -264,7 +265,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
 
 
 def test_add_devices_from_module(empty_context: BlueskyContext):
@@ -720,7 +721,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):
@@ -791,7 +792,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):