Reimplement DeviceCollector as auto_init_devices

Functionality is the same, but the public interface is now a function rather than a class.
Also makes sure that Devices that exist on entry, and are redeclared in the context manager are found.

docs/examples/epics_demo.py

@@ -7,7 +7,7 @@
 from bluesky.utils import ProgressBarManager, register_transform
 from ophyd import Component, Device, EpicsSignal, EpicsSignalRO
 
-from ophyd_async.core import DeviceCollector
+from ophyd_async.core import auto_init_devices
 from ophyd_async.epics import demo
 
 # Create a run engine, with plotting, progressbar and transform
@@ -31,7 +31,7 @@ class OldSensor(Device):
 det_old = OldSensor(pv_prefix, name="det_old")
 
 # Create ophyd-async devices
-with DeviceCollector():
+with auto_init_devices():
     det = demo.Sensor(pv_prefix)
     det_group = demo.SensorGroup(pv_prefix)
     samp = demo.SampleStage(pv_prefix)

docs/explanations/event-loop-choice.rst

@@ -1,7 +1,7 @@
 Device Collector Event-Loop Choice
 ----------------------------------
 
-In a sync context, the ophyd-async :python:`DeviceCollector` requires the bluesky event-loop
+In a sync context, the ophyd-async :python:`auto_init_devices` requires the bluesky event-loop
 to connect to devices. In an async context, it does not.
 
 Sync Context
@@ -14,25 +14,25 @@ The following will fail if :python:`RE = RunEngine()` has not been called alread
 
 .. code:: python
 
-  with DeviceCollector():
+  with auto_init_devices():
       device1 = Device1(prefix)
       device2 = Device2(prefix)
       device3 = Device3(prefix)
 
-The :python:`DeviceCollector` connects to devices in the event-loop created in the run-engine.
+The :python:`auto_init_devices` connects to devices in the event-loop created in the run-engine.
 
 
 Async Context
 =============
 
 In an async context device connection is decoupled from the run-engine.
-The following attempts connection to all the devices in the :python:`DeviceCollector`
+The following attempts connection to all the devices in the :python:`auto_init_devices`
 before or after run-engine initialization.
 
 .. code:: python
 
   async def connection_function() :
-      async with DeviceCollector():
+      async with auto_init_devices():
           device1 = Device1(prefix)
           device2 = Device2(prefix)
           device3 = Device3(prefix)
@@ -43,8 +43,8 @@ The devices will be unable to be used in the run-engine unless they share the sa
 When the run-engine is initialised it will create a new background event-loop to use if one
 is not passed in with :python:`RunEngine(loop=loop)`.
 
-If the user wants to use devices in the async :python:`DeviceCollector` within the run-engine
+If the user wants to use devices in the async :python:`auto_init_devices` within the run-engine
 they can either:
 
-* Run the :python:`DeviceCollector` first and pass the event-loop into the run-engine.
-* Initialize the run-engine first and run the :python:`DeviceCollector` using the bluesky event-loop.
+* Run the :python:`auto_init_devices` first and pass the event-loop into the run-engine.
+* Initialize the run-engine first and run the :python:`auto_init_devices` using the bluesky event-loop.

docs/how-to/write-tests-for-devices.rst

@@ -24,7 +24,7 @@ Async Tests
 Mock Backend
 ------------
 
-Ophyd devices initialized with a mock backend behave in a similar way to mocks, without requiring you to mock out all the dependencies and internals. The `DeviceCollector` can initialize any number of devices, and their signals and sub-devices (recursively), with a mock backend.
+Ophyd devices initialized with a mock backend behave in a similar way to mocks, without requiring you to mock out all the dependencies and internals. The `auto_init_devices` can initialize any number of devices, and their signals and sub-devices (recursively), with a mock backend.
 
 .. literalinclude:: ../../tests/epics/demo/test_demo.py
    :pyobject: mock_sensor

docs/tutorials/using-existing-devices.rst

@@ -47,7 +47,7 @@ Finally we create the Ophyd Async devices imported from the `epics.demo` module:
     :language: python
     :start-after: # Create ophyd-async devices
 
-The first thing to note is `with`. This uses a `DeviceCollector` as a context
+The first thing to note is `with`. This uses `auto_init_devices` as a context
 manager to collect up the top level `Device` instances created in the context,
 and run the following:
 

src/ophyd_async/core/__init__.py

@@ -5,7 +5,7 @@
     StandardDetector,
     TriggerInfo,
 )
-from ._device import Device, DeviceCollector, DeviceConnector, DeviceVector
+from ._device import Device, DeviceConnector, DeviceVector, auto_init_devices
 from ._device_filler import DeviceFiller
 from ._flyer import FlyerController, StandardFlyer
 from ._hdf_dataset import HDFDataset, HDFFile
@@ -87,7 +87,7 @@
     "TriggerInfo",
     "Device",
     "DeviceConnector",
-    "DeviceCollector",
+    "auto_init_devices",
     "DeviceVector",
     "DeviceFiller",
     "StandardFlyer",

src/ophyd_async/core/_device.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import sys
-from collections.abc import Coroutine, Iterator, Mapping, MutableMapping
+from collections.abc import Awaitable, Callable, Iterator, Mapping, MutableMapping
 from functools import cached_property
 from logging import LoggerAdapter, getLogger
 from typing import Any, TypeVar
@@ -254,54 +254,18 @@ def __hash__(self):  # to allow DeviceVector to be used as dict keys and in sets
         return hash(id(self))
 
 
-class DeviceCollector:
-    """Collector of top level Device instances to be used as a context manager
-
-    Parameters
-    ----------
-    set_name:
-        If True, call ``device.set_name(variable_name)`` on all collected
-        Devices
-    child_name_separator:
-        Use this as a separator if we call ``set_name``.
-    connect:
-        If True, call ``device.connect(mock)`` in parallel on all
-        collected Devices
-    mock:
-        If True, connect Signals in simulation mode
-    timeout:
-        How long to wait for connect before logging an exception
-
-    Notes
-    -----
-    Example usage::
-
-        [async] with DeviceCollector():
-            t1x = motor.Motor("BLxxI-MO-TABLE-01:X")
-            t1y = motor.Motor("pva://BLxxI-MO-TABLE-01:Y")
-            # Names and connects devices here
-        assert t1x.comm.velocity.source
-        assert t1x.name == "t1x"
+class DeviceContextManager:
+    """Sync/Async Context Manager that finds all the Devices declared within it.
 
+    Used in `auto_init`
     """
 
-    def __init__(
-        self,
-        set_name=True,
-        child_name_separator: str = "-",
-        connect=True,
-        mock=False,
-        timeout: float = 10.0,
-    ):
-        self._set_name = set_name
-        self._child_name_separator = child_name_separator
-        self._connect = connect
-        self._mock = mock
-        self._timeout = timeout
-        self._names_on_enter: set[str] = set()
-        self._objects_on_exit: dict[str, Any] = {}
-
-    def _caller_locals(self):
+    def __init__(self, process_devices: Callable[[dict[str, Device]], Awaitable[None]]):
+        self._process_devices = process_devices
+        self._locals_on_enter: dict[str, Any] = {}
+        self._locals_on_exit: dict[str, Any] = {}
+
+    def _caller_locals(self) -> dict[str, Any]:
         """Walk up until we find a stack frame that doesn't have us as self"""
         try:
             raise ValueError
@@ -314,34 +278,18 @@ def _caller_locals(self):
                 assert (
                     caller_frame
                 ), "No previous frame to the one with self in it, this shouldn't happen"
-            return caller_frame.f_locals
+            return caller_frame.f_locals.copy()
 
-    def __enter__(self) -> DeviceCollector:
+    def __enter__(self) -> DeviceContextManager:
         # Stash the names that were defined before we were called
-        self._names_on_enter = set(self._caller_locals())
+        self._locals_on_enter = self._caller_locals()
         return self
 
-    async def __aenter__(self) -> DeviceCollector:
+    async def __aenter__(self) -> DeviceContextManager:
         return self.__enter__()
 
-    async def _on_exit(self) -> None:
-        # Name and kick off connect for devices
-        connect_coroutines: dict[str, Coroutine] = {}
-        for name, obj in self._objects_on_exit.items():
-            if name not in self._names_on_enter and isinstance(obj, Device):
-                if self._set_name and not obj.name:
-                    obj.set_name(name, child_name_separator=self._child_name_separator)
-                if self._connect:
-                    connect_coroutines[name] = obj.connect(
-                        self._mock, timeout=self._timeout
-                    )
-
-        # Connect to all the devices
-        if connect_coroutines:
-            await wait_for_connection(**connect_coroutines)
-
     async def __aexit__(self, type, value, traceback):
-        self._objects_on_exit = self._caller_locals()
+        self._locals_on_exit = self._caller_locals()
         await self._on_exit()
 
     def __exit__(self, type_, value, traceback):
@@ -350,7 +298,7 @@ def __exit__(self, type_, value, traceback):
                 "Cannot use DeviceConnector inside a plan, instead use "
                 "`yield from ophyd_async.plan_stubs.ensure_connected(device)`"
             )
-        self._objects_on_exit = self._caller_locals()
+        self._locals_on_exit = self._caller_locals()
         try:
             fut = call_in_bluesky_event_loop(self._on_exit())
         except RuntimeError as e:
@@ -360,3 +308,62 @@ def __exit__(self, type_, value, traceback):
                 "user/explanations/event-loop-choice.html for more info."
             ) from e
         return fut
+
+    async def _on_exit(self) -> None:
+        # Find all the devices
+        devices = {
+            name: obj
+            for name, obj in self._locals_on_exit.items()
+            if isinstance(obj, Device) and self._locals_on_enter.get(name) is not obj
+        }
+        # Call the provided process function on them
+        await self._process_devices(devices)
+
+
+def auto_init_devices(
+    set_name=True,
+    child_name_separator: str = "-",
+    connect=True,
+    mock=False,
+    timeout: float = 10.0,
+) -> DeviceContextManager:
+    """Auto initialise top level Device instances to be used as a context manager
+
+    Parameters
+    ----------
+    set_name:
+        If True, call ``device.set_name(variable_name)`` on all Devices
+        created within the context manager that have an empty ``name``
+    child_name_separator:
+        Use this as a separator if we call ``set_name``.
+    connect:
+        If True, call ``device.connect(mock, timeout)`` in parallel on all
+        Devices created within the context manager
+    mock:
+        If True, connect Signals in mock mode
+    timeout:
+        How long to wait for connect before logging an exception
+
+    Notes
+    -----
+    Example usage::
+
+        [async] with auto_init_devices():
+            t1x = motor.Motor("BLxxI-MO-TABLE-01:X")
+            t1y = motor.Motor("pva://BLxxI-MO-TABLE-01:Y")
+            # Names and connects devices here
+        assert t1x.name == "t1x"
+    """
+
+    async def process_devices(devices: dict[str, Device]):
+        if set_name:
+            for name, device in devices.items():
+                if not device.name:
+                    device.set_name(name, child_name_separator=child_name_separator)
+        if connect:
+            coros = {
+                name: device.connect(mock, timeout) for name, device in devices.items()
+            }
+            await wait_for_connection(**coros)
+
+    return DeviceContextManager(process_devices)

system_tests/epics/eiger/test_eiger_system.py

@@ -9,8 +9,8 @@
 from ophyd_async.core import (
     DetectorTrigger,
     Device,
-    DeviceCollector,
     StaticPathProvider,
+    auto_init_devices,
 )
 from ophyd_async.epics.core import epics_signal_rw
 from ophyd_async.epics.eiger import EigerDetector, EigerTriggerInfo
@@ -47,7 +47,7 @@ def RE():
 
 @pytest.fixture
 async def setup_device(RE, ioc_prefixes):
-    async with DeviceCollector():
+    async with auto_init_devices():
         device = SetupDevice(ioc_prefixes[0], ioc_prefixes[1] + "FP:")
     await asyncio.gather(
         device.header_detail.set("all"),
@@ -62,7 +62,7 @@ async def setup_device(RE, ioc_prefixes):
 @pytest.fixture
 async def test_eiger(RE, ioc_prefixes) -> EigerDetector:
     provider = StaticPathProvider(lambda: "test_eiger", Path(SAVE_PATH))
-    async with DeviceCollector():
+    async with auto_init_devices():
         test_eiger = EigerDetector("", provider, ioc_prefixes[0], ioc_prefixes[1])
 
     return test_eiger