feat: Add depth parameter to get_devices

Author: Ware, Joseph (DLSLtd,RAL,LSCI)

docs/reference/openapi.yaml

@@ -340,7 +340,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 +361,32 @@ paths:
     get:
       description: Retrieve information about all available devices.
       operationId: get_devices_devices_get
+      parameters:
+      - description: Maximum depth of children to return
+        in: query
+        name: depth
+        required: false
+        schema:
+          anyOf:
+          - type: integer
+          - const: all
+            type: string
+          default: 0
+          ge: 0
+          title: Depth
       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

@@ -6,6 +6,7 @@
 from functools import wraps
 from pathlib import Path
 from pprint import pprint
+from typing import Literal
 
 import click
 from bluesky.callbacks.best_effort import BestEffortCallback
@@ -165,11 +166,18 @@ def get_plans(obj: dict) -> None:
 
 @controller.command(name="devices")
 @check_connection
+@click.option(
+    "-d",
+    "--depth",
+    type=click.IntRange(min=0) | Literal["all"],
+    required=False,
+    help="Maximum depth of children to return",
+)
 @click.pass_obj
-def get_devices(obj: dict) -> None:
+def get_devices(obj: dict, depth: int | Literal["all"] = 0) -> 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(depth))
 
 
 @controller.command(name="listen")

src/blueapi/client/client.py

@@ -1,5 +1,6 @@
 import time
 from concurrent.futures import Future
+from typing import Literal
 
 from bluesky_stomp.messaging import MessageContext, StompClient
 from bluesky_stomp.models import Broker
@@ -91,15 +92,15 @@ def get_plan(self, name: str) -> PlanModel:
         return self._rest.get_plan(name)
 
     @start_as_current_span(TRACER)
-    def get_devices(self) -> DeviceResponse:
+    def get_devices(self, depth: int | Literal["all"]) -> DeviceResponse:
         """
         List devices available
 
         Returns:
             DeviceResponse: Devices that can be used in plans
         """
 
-        return self._rest.get_devices()
+        return self._rest.get_devices(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, depth: int | Literal["all"]) -> DeviceResponse:
+        return self._request_and_deserialize(
+            "/devices", DeviceResponse, params={"depth": depth}
+        )
 
     def get_device(self, name: str) -> DeviceModel:
         return self._request_and_deserialize(f"/devices/{name}", DeviceModel)

src/blueapi/service/interface.py

@@ -1,7 +1,7 @@
 import logging
 from collections.abc import Mapping
 from functools import cache
-from typing import Any
+from typing import Any, Literal
 
 from bluesky_stomp.messaging import StompClient
 from bluesky_stomp.models import Broker, DestinationBase, MessageTopic
@@ -176,9 +176,14 @@ def get_plan(name: str) -> PlanModel:
     return PlanModel.from_plan(context().plans[name])
 
 
-def get_devices() -> list[DeviceModel]:
+def get_devices(depth: int | Literal["all"]) -> 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, depth)
+        if model.protocols
+    ]
 
 
 def get_device(name: str) -> DeviceModel:

src/blueapi/service/main.py

@@ -1,6 +1,6 @@
 from collections.abc import Awaitable, Callable
 from contextlib import asynccontextmanager
-from typing import Annotated
+from typing import Annotated, Literal
 
 import jwt
 from fastapi import (
@@ -10,6 +10,7 @@
     Depends,
     FastAPI,
     HTTPException,
+    Query,
     Request,
     Response,
     status,
@@ -53,7 +54,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
 
@@ -228,9 +229,18 @@ def get_plan_by_name(
 @start_as_current_span(TRACER)
 def get_devices(
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
+    depth: Annotated[
+        int | Literal["all"],
+        Query(
+            description="Maximum depth of children to return",
+            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, depth)
     return DeviceResponse(devices=devices)
 
 

src/blueapi/service/model.py

@@ -1,9 +1,11 @@
 import uuid
 from collections.abc import Iterable
 from enum import Enum
-from typing import Annotated, Any
+from typing import Annotated, Any, Literal
 
 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
 
@@ -40,10 +42,37 @@ 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)))
 
+    @classmethod
+    def from_device_tree(
+        cls, root: Device, max_depth: int | Literal["all"]
+    ) -> list["DeviceModel"]:
+        if max_depth == 0:
+            return [cls.from_device(root)]
+        if isinstance(root, AsyncDevice):
+            # Breadth-first iteration through child devices, stopping at max_depth
+            async_devices: list[AsyncDevice] = [root]
+            branches: list[AsyncDevice] = [child[1] for child in root.children()]
+            depth = 0
+            while (max_depth == "all" or depth < max_depth) and branches:
+                async_devices += branches
+                branches = [
+                    child[1] for branch in branches for child in branch.children()
+                ]
+                depth += 1
+            return [cls.from_device(device) for device in async_devices]
+        elif isinstance(root, SyncDevice):
+            sync_devices: list[SyncDevice] = [root]
+            # Depth-first iteration through components, keeping any below max_depth
+            for component in root.walk_components():
+                if max_depth == "all" or len(component.ancestors) <= max_depth:
+                    sync_devices.append(component.item)
+            return [cls.from_device(device) for device in sync_devices]
+        return [cls.from_device(root)]
+
 
 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(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

@@ -1,5 +1,6 @@
 import uuid
 from collections.abc import Callable
+from typing import Literal
 from unittest.mock import MagicMock, Mock, call, patch
 
 import pytest
@@ -39,6 +40,21 @@
         DeviceModel(name="bar", protocols=[]),
     ]
 )
+DEVICES_AND_CHILDREN = DeviceResponse(
+    devices=[
+        DeviceModel(name="foo", protocols=[]),
+        DeviceModel(name="bar", protocols=[]),
+        DeviceModel(name="foo-bar", protocols=[]),
+    ]
+)
+DEVICES_AND_ALL_DESCENDENTS = DeviceResponse(
+    devices=[
+        DeviceModel(name="foo", protocols=[]),
+        DeviceModel(name="bar", protocols=[]),
+        DeviceModel(name="foo-bar", protocols=[]),
+        DeviceModel(name="foo-bar-baz", protocols=[]),
+    ]
+)
 DEVICE = DeviceModel(name="foo", protocols=[])
 TASK = TrackableTask(task_id="foo", task=Task(name="bar", params={}))
 TASKS = TasksListResponse(tasks=[TASK])
@@ -67,11 +83,18 @@
 
 @pytest.fixture
 def mock_rest() -> BlueapiRestClient:
+    def get_devices(depth: int | Literal["all"]) -> DeviceResponse:
+        if depth == "all" or depth > 1:
+            return DEVICES_AND_ALL_DESCENDENTS
+        if 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 +144,7 @@ def test_get_nonexistant_plan(
 
 
 def test_get_devices(client: BlueapiClient):
-    assert client.get_devices() == DEVICES
+    assert client.get_devices(depth=0) == DEVICES
 
 
 def test_get_device(client: BlueapiClient):
@@ -511,7 +534,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(depth=0)
 
 
 def test_get_device_span_ok(exporter: JsonObjectSpanExporter, client: BlueapiClient):