Add children method to Device and DeviceVector

src/ophyd_async/core/__init__.py

@@ -13,7 +13,7 @@
     set_sim_value,
     wait_for_value,
 )
-from ._device.device import Device, connect_children, get_device_children, name_children
+from ._device.device import Device
 from ._device.device_collector import DeviceCollector
 from ._device.device_vector import DeviceVector
 from ._device.standard_readable import StandardReadable
@@ -45,9 +45,6 @@
     "set_sim_value",
     "wait_for_value",
     "Device",
-    "connect_children",
-    "get_device_children",
-    "name_children",
     "DeviceCollector",
     "DeviceVector",
     "StandardReadable",

src/ophyd_async/core/_device/device.py

@@ -1,6 +1,7 @@
 """Base device"""
+from __future__ import annotations
 
-from typing import Generator, Optional, Tuple
+from typing import Iterator, Optional, Tuple
 
 from bluesky.protocols import HasName
 
@@ -15,7 +16,7 @@ class Device(HasName):
 
     _name: str = ""
     #: The parent Device if it exists
-    parent: Optional["Device"] = None
+    parent: Optional[Device] = None
 
     def __init__(self, name: str = "") -> None:
         self.set_name(name)
@@ -24,6 +25,11 @@ def __init__(self, name: str = "") -> None:
     def name(self) -> str:
         """Return the name of the Device"""
         return self._name
+
+    def children(self) -> Iterator[Tuple[str, Device]]:
+        for attr_name, attr in self.__dict__.items():
+            if attr_name != "parent" and isinstance(attr, Device):
+                yield attr_name, attr
 
     def set_name(self, name: str):
         """Set ``self.name=name`` and each ``self.child.name=name+"-child"``.
@@ -34,7 +40,10 @@ def set_name(self, name: str):
             New name to set
         """
         self._name = name
-        name_children(self, name)
+        for attr_name, child in self.children():
+            child_name = f"{name}-{attr_name.rstrip('_')}" if name else ""
+            child.set_name(child_name)
+            child.parent = self
 
     async def connect(self, sim: bool = False):
         """Connect self and all child Devices.
@@ -44,35 +53,8 @@ async def connect(self, sim: bool = False):
         sim:
             If True then connect in simulation mode.
         """
-        await connect_children(self, sim)
-
-
-async def connect_children(device: Device, sim: bool):
-    """Call ``child.connect(sim)`` on all child devices in parallel.
-
-    Typically used to implement `Device.connect` like this::
-
-        async def connect(self, sim=False):
-            await connect_children(self, sim)
-    """
-
-    coros = {
-        name: child_device.connect(sim)
-        for name, child_device in get_device_children(device)
-    }
-    if coros:
-        await wait_for_connection(**coros)
-
-
-def name_children(device: Device, name: str):
-    """Call ``child.set_name(child_name)`` on all child devices in series."""
-    for attr_name, child in get_device_children(device):
-        child_name = f"{name}-{attr_name.rstrip('_')}" if name else ""
-        child.set_name(child_name)
-        child.parent = device
-
-
-def get_device_children(device: Device) -> Generator[Tuple[str, Device], None, None]:
-    for attr_name, attr in device.__dict__.items():
-        if attr_name != "parent" and isinstance(attr, Device):
-            yield attr_name, attr
+        coros = {
+            name: child_device.connect(sim) for name, child_device in self.children()
+        }
+        if coros:
+            await wait_for_connection(**coros)

src/ophyd_async/core/_device/device_vector.py

@@ -1,6 +1,6 @@
 """Dictionary which can contain mappings between integers and devices."""
 
-from typing import Dict, TypeVar
+from typing import Dict, Generator, Tuple, TypeVar
 
 from ..utils import wait_for_connection
 from .device import Device
@@ -9,12 +9,7 @@
 
 
 class DeviceVector(Dict[int, VT], Device):
-    def set_name(self, parent_name: str):
-        self._name = parent_name
-        for name, device in self.items():
-            device.set_name(f"{parent_name}-{name}")
-            device.parent = self
-
-    async def connect(self, sim: bool = False):
-        coros = {str(k): d.connect(sim) for k, d in self.items()}
-        await wait_for_connection(**coros)
+    def children(self) -> Generator[Tuple[str, Device], None, None]:
+        for attr_name, attr in self.items():
+            if isinstance(attr, Device):
+                yield str(attr_name), attr

tests/core/_device/test_device.py

@@ -3,13 +3,7 @@
 
 import pytest
 
-from ophyd_async.core import (
-    Device,
-    DeviceCollector,
-    DeviceVector,
-    get_device_children,
-    wait_for_connection,
-)
+from ophyd_async.core import Device, DeviceCollector, DeviceVector, wait_for_connection
 
 
 class DummyBaseDevice(Device):
@@ -35,15 +29,25 @@ def parent() -> DummyDeviceGroup:
     return DummyDeviceGroup("parent")
 
 
-def test_get_device_children(parent: DummyDeviceGroup):
+def test_device_children(parent: DummyDeviceGroup):
     names = ["child1", "child2", "dict_with_children"]
-    for idx, (name, child) in enumerate(get_device_children(parent)):
+    for idx, (name, child) in enumerate(parent.children()):
         assert name == names[idx]
         assert (
             type(child) is DummyBaseDevice
             if name.startswith("child")
             else type(child) is DeviceVector
         )
+        assert child.parent == parent
+
+
+def test_device_vector_children():
+    parent = DummyDeviceGroup("root")
+
+    device_vector_children = [
+        (name, child) for name, child in parent.dict_with_children.children()
+    ]
+    assert device_vector_children == [("123", parent.dict_with_children[123])]
 
 
 async def test_children_of_device_have_set_names_and_get_connected(