Add depth parameter to get_devices

Author: DiamondJoseph

docs/reference/openapi.yaml

@@ -4,6 +4,9 @@ components:
       additionalProperties: false
       description: Representation of a device
       properties:
+        address:
+          title: Address
+          type: string
         name:
           description: Name of the device
           title: Name
@@ -17,6 +20,7 @@ components:
       required:
       - name
       - protocols
+      - address
       title: DeviceModel
       type: object
     DeviceResponse:
@@ -340,7 +344,7 @@ components:
       type: object
 info:
   title: BlueAPI Control
-  version: 0.0.10
+  version: 0.0.11
 openapi: 3.1.0
 paths:
   /config/oidc:
@@ -361,13 +365,29 @@ paths:
     get:
       description: Retrieve information about all available devices.
       operationId: get_devices_devices_get
+      parameters:
+      - description: Maximum depth of children to return, -1 for all
+        in: query
+        name: max_depth
+        required: false
+        schema:
+          default: 0
+          minimum: 0
+          title: Max Depth
+          type: integer
       responses:
         '200':
           content:
             application/json:
               schema:
                 $ref: '#/components/schemas/DeviceResponse'
           description: Successful Response
+        '422':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+          description: Validation Error
       summary: Get Devices
   /devices/{name}:
     get:

src/blueapi/cli/cli.py

@@ -164,11 +164,19 @@ def get_plans(obj: dict) -> None:
 
 @controller.command(name="devices")
 @check_connection
+@click.option(
+    "-d",
+    "--max_depth",
+    type=click.IntRange(-1),
+    required=False,
+    help="Maximum depth of children to return: -1 for all",
+    default=0,
+)
 @click.pass_obj
-def get_devices(obj: dict) -> None:
+def get_devices(obj: dict, max_depth: int) -> None:
     """Get a list of devices available for the worker to use"""
     client: BlueapiClient = obj["client"]
-    obj["fmt"].display(client.get_devices())
+    obj["fmt"].display(client.get_devices(max_depth))
 
 
 @controller.command(name="listen")

src/blueapi/cli/format.py

@@ -12,6 +12,7 @@
 
 from blueapi.core.bluesky_types import DataEvent
 from blueapi.service.model import (
+    DeviceModel,
     DeviceResponse,
     PlanResponse,
     PythonEnvironmentResponse,
@@ -62,7 +63,7 @@ def display_full(obj: Any, stream: Stream):
                     print(indent(json.dumps(schema, indent=2), "        "))
         case DeviceResponse(devices=devices):
             for dev in devices:
-                print(dev.name)
+                print(_format_name(dev))
                 for proto in dev.protocols:
                     print(f"    {proto}")
         case DataEvent(name=name, doc=doc):
@@ -124,7 +125,7 @@ def display_compact(obj: Any, stream: Stream):
                         print(f"      {arg}={_describe_type(spec, req)}")
         case DeviceResponse(devices=devices):
             for dev in devices:
-                print(dev.name)
+                print(_format_name(dev))
                 print(
                     indent(
                         textwrap.fill(
@@ -164,6 +165,12 @@ def display_compact(obj: Any, stream: Stream):
             FALLBACK(other, stream=stream)
 
 
+def _format_name(device: DeviceModel) -> str:
+    if not device.address or device.address == device.name:
+        return device.name
+    return f"{device.name} @ {device.address}"
+
+
 def _describe_type(spec: dict[Any, Any], required: bool = False):
     disp = ""
     match spec.get("type"):

src/blueapi/client/client.py

@@ -90,16 +90,16 @@ def get_plan(self, name: str) -> PlanModel:
         """
         return self._rest.get_plan(name)
 
-    @start_as_current_span(TRACER)
-    def get_devices(self) -> DeviceResponse:
+    @start_as_current_span(TRACER, "max_depth")
+    def get_devices(self, max_depth: int) -> DeviceResponse:
         """
         List devices available
 
         Returns:
             DeviceResponse: Devices that can be used in plans
         """
 
-        return self._rest.get_devices()
+        return self._rest.get_devices(max_depth)
 
     @start_as_current_span(TRACER, "name")
     def get_device(self, name: str) -> DeviceModel:

src/blueapi/client/rest.py

@@ -71,8 +71,10 @@ def get_plans(self) -> PlanResponse:
     def get_plan(self, name: str) -> PlanModel:
         return self._request_and_deserialize(f"/plans/{name}", PlanModel)
 
-    def get_devices(self) -> DeviceResponse:
-        return self._request_and_deserialize("/devices", DeviceResponse)
+    def get_devices(self, max_depth: int) -> DeviceResponse:
+        return self._request_and_deserialize(
+            "/devices", DeviceResponse, params={"max_depth": max_depth}
+        )
 
     def get_device(self, name: str) -> DeviceModel:
         return self._request_and_deserialize(f"/devices/{name}", DeviceModel)

src/blueapi/core/device_lookup.py

@@ -1,5 +1,7 @@
 from typing import Any
 
+from ophyd_async.core import DeviceVector
+
 from .bluesky_types import Device, is_bluesky_compatible_device
 
 
@@ -28,6 +30,8 @@ def find_component(obj: Any, addr: list[str]) -> Device | None:
     # Otherwise, we error.
     if isinstance(obj, dict):
         component = obj.get(head)
+    elif isinstance(obj, DeviceVector):
+        component = obj.get(int(head))
     elif is_bluesky_compatible_device(obj):
         component = getattr(obj, head, None)
     else:

src/blueapi/service/interface.py

@@ -177,9 +177,14 @@ def get_plan(name: str) -> PlanModel:
     return PlanModel.from_plan(context().plans[name])
 
 
-def get_devices() -> list[DeviceModel]:
+def get_devices(max_depth: int) -> list[DeviceModel]:
     """Get all available devices in the BlueskyContext"""
-    return [DeviceModel.from_device(device) for device in context().devices.values()]
+    return [
+        model
+        for device in context().devices.values()
+        for model in DeviceModel.from_device_tree(device, max_depth)
+        if model.protocols
+    ]
 
 
 def get_device(name: str) -> DeviceModel:

src/blueapi/service/main.py

@@ -11,6 +11,7 @@
     Depends,
     FastAPI,
     HTTPException,
+    Query,
     Request,
     Response,
     status,
@@ -54,7 +55,7 @@
 from .runner import WorkerDispatcher
 
 #: API version to publish in OpenAPI schema
-REST_API_VERSION = "0.0.10"
+REST_API_VERSION = "0.0.11"
 
 RUNNER: WorkerDispatcher | None = None
 
@@ -231,9 +232,18 @@ def get_plan_by_name(
 @start_as_current_span(TRACER)
 def get_devices(
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
+    max_depth: Annotated[
+        int,
+        Query(
+            description="Maximum depth of children to return, -1 for all",
+            ge=0,
+            # https://github.com/fastapi/fastapi/discussions/13473
+            json_schema_extra={"description": None},
+        ),
+    ] = 0,
 ) -> DeviceResponse:
     """Retrieve information about all available devices."""
-    devices = runner.run(interface.get_devices)
+    devices = runner.run(interface.get_devices, max_depth)
     return DeviceResponse(devices=devices)
 
 

src/blueapi/service/model.py

@@ -4,6 +4,8 @@
 from typing import Annotated, Any
 
 from bluesky.protocols import HasName
+from ophyd import Device as SyncDevice
+from ophyd_async.core import Device as AsyncDevice
 from pydantic import Field
 from pydantic.json_schema import SkipJsonSchema
 
@@ -34,16 +36,69 @@ class DeviceModel(BlueapiBaseModel):
     protocols: list[ProtocolInfo] = Field(
         description="Protocols that a device conforms to, indicating its capabilities"
     )
+    address: str
 
     @classmethod
     def from_device(cls, device: Device) -> "DeviceModel":
         name = device.name if isinstance(device, HasName) else _UNKNOWN_NAME
-        return cls(name=name, protocols=list(_protocol_info(device)))
+        return cls(name=name, protocols=list(_protocol_info(device)), address=name)
+
+    @classmethod
+    def from_device_tree(cls, root: Device, max_depth: int) -> list["DeviceModel"]:
+        if isinstance(root, AsyncDevice):
+            return [
+                DeviceModel(
+                    name=device.name,
+                    protocols=list(_protocol_info(device)),
+                    address=address,
+                )
+                for address, device in _from_async_device(
+                    root, max_depth=max_depth
+                ).items()
+            ]
+        if isinstance(root, SyncDevice):
+            return [
+                DeviceModel(
+                    name=device.name,
+                    protocols=list(_protocol_info(device)),
+                    address=address,
+                )
+                for address, device in _from_sync_device(
+                    root, max_depth=max_depth
+                ).items()
+            ]
+        return [DeviceModel.from_device(root)]
+
+
+def _from_async_device(root: AsyncDevice, max_depth: int) -> dict[str, AsyncDevice]:
+    depth = 0
+    devices: dict[str, AsyncDevice] = {root.name: root}
+    branches: dict[str, AsyncDevice] = {root.name: root}
+    while branches and (max_depth == -1 or depth < max_depth):
+        leaves: dict[str, AsyncDevice] = {}
+        for addr, parent in branches.items():
+            for suffix, child in parent.children():
+                leaves[f"{addr}.{suffix}"] = child
+        devices.update(leaves)
+        branches = leaves
+        depth += 1
+    return devices
+
+
+def _from_sync_device(root: SyncDevice, max_depth: int) -> dict[str, SyncDevice]:
+    return {
+        root.name: root,
+        **{
+            k.dotted_name: k.item
+            for k in root.walk_signals()
+            if max_depth == -1 or len(k.ancestors) <= max_depth
+        },
+    }
 
 
 def _protocol_info(device: Device) -> Iterable[ProtocolInfo]:
     for protocol in BLUESKY_PROTOCOLS:
-        if isinstance(device, protocol):
+        if isinstance(device, protocol) and protocol is not AsyncDevice:
             yield ProtocolInfo(
                 name=protocol.__name__,
                 types=[arg.__name__ for arg in generic_bounds(device, protocol)],

tests/system_tests/test_blueapi_system.py

@@ -162,7 +162,7 @@ def test_get_non_existent_plan(client: BlueapiClient):
 
 
 def test_get_devices(client: BlueapiClient, expected_devices: DeviceResponse):
-    retrieved_devices = client.get_devices()
+    retrieved_devices = client.get_devices(max_depth=0)
     retrieved_devices.devices.sort(key=lambda x: x.name)
     expected_devices.devices.sort(key=lambda x: x.name)
 

tests/unit_tests/client/test_client.py

@@ -35,11 +35,26 @@
 PLAN = PlanModel(name="foo")
 DEVICES = DeviceResponse(
     devices=[
-        DeviceModel(name="foo", protocols=[]),
-        DeviceModel(name="bar", protocols=[]),
+        DeviceModel(name="foo", protocols=[], address="foo"),
+        DeviceModel(name="bar", protocols=[], address="bar"),
     ]
 )
-DEVICE = DeviceModel(name="foo", protocols=[])
+DEVICES_AND_CHILDREN = DeviceResponse(
+    devices=[
+        DeviceModel(name="foo", protocols=[], address="foo"),
+        DeviceModel(name="bar", protocols=[], address="bar"),
+        DeviceModel(name="foo-bar", protocols=[], address="foo.bar"),
+    ]
+)
+DEVICES_AND_ALL_DESCENDENTS = DeviceResponse(
+    devices=[
+        DeviceModel(name="foo", protocols=[], address="foo"),
+        DeviceModel(name="bar", protocols=[], address="bar"),
+        DeviceModel(name="foo-bar", protocols=[], address="foo.bar"),
+        DeviceModel(name="foo-bar-baz", protocols=[], address="foo.bar.baz"),
+    ]
+)
+DEVICE = DeviceModel(name="foo", protocols=[], address="foo")
 TASK = TrackableTask(task_id="foo", task=Task(name="bar", params={}))
 TASKS = TasksListResponse(tasks=[TASK])
 ACTIVE_TASK = WorkerTask(task_id="bar")
@@ -67,11 +82,18 @@
 
 @pytest.fixture
 def mock_rest() -> BlueapiRestClient:
+    def get_devices(max_depth: int) -> DeviceResponse:
+        if max_depth == -1 or max_depth > 1:
+            return DEVICES_AND_ALL_DESCENDENTS
+        if max_depth == 1:
+            return DEVICES_AND_CHILDREN
+        return DEVICES
+
     mock = Mock(spec=BlueapiRestClient)
 
     mock.get_plans.return_value = PLANS
     mock.get_plan.return_value = PLAN
-    mock.get_devices.return_value = DEVICES
+    mock.get_devices.side_effect = get_devices
     mock.get_device.return_value = DEVICE
     mock.get_state.return_value = WorkerState.IDLE
     mock.get_task.return_value = TASK
@@ -121,7 +143,7 @@ def test_get_nonexistant_plan(
 
 
 def test_get_devices(client: BlueapiClient):
-    assert client.get_devices() == DEVICES
+    assert client.get_devices(max_depth=0) == DEVICES
 
 
 def test_get_device(client: BlueapiClient):
@@ -511,7 +533,7 @@ def test_get_plan_span_ok(exporter: JsonObjectSpanExporter, client: BlueapiClien
 
 def test_get_devices_span_ok(exporter: JsonObjectSpanExporter, client: BlueapiClient):
     with asserting_span_exporter(exporter, "get_devices"):
-        client.get_devices()
+        client.get_devices(max_depth=0)
 
 
 def test_get_device_span_ok(exporter: JsonObjectSpanExporter, client: BlueapiClient):