ophyd-async 0.9.0a1__py3-none-any.whl → 0.10.0a1__py3-none-any.whl

ophyd_async/core/_device.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 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
@@ -17,7 +17,7 @@ class DeviceConnector:
     """Defines how a `Device` should be connected and type hints processed."""
 
     def create_children_from_annotations(self, device: Device):
-        """Used when children can be created from introspecting the hardware.
+        """Use when children can be created from introspecting the hardware.
 
         Some control systems allow introspection of a device to determine what
         children it has. To allow this to work nicely with typing we add these
@@ -26,15 +26,20 @@ class DeviceConnector:
             my_signal: SignalRW[int]
             my_device: MyDevice
 
-        This method will be run during ``Device.__init__``, and is responsible
+        This method will be run during `Device.__init__`, and is responsible
         for turning all of those type hints into real Signal and Device instances.
 
         Subsequent runs of this function should do nothing, to allow it to be
         called early in Devices that need to pass references to their children
-        during ``__init__``.
+        during `__init__`.
         """
 
     async def connect_mock(self, device: Device, mock: LazyMock):
+        """Use during [](#Device.connect) with `mock=True`.
+
+        This is called when there is no cached connect done in `mock=True`
+        mode. It connects the Device and all its children in mock mode.
+        """
         # Connect serially, no errors to gather up as in mock mode
         exceptions: dict[str, Exception] = {}
         for name, child_device in device.children():
@@ -46,11 +51,10 @@ class DeviceConnector:
             raise NotConnected.with_other_exceptions_logged(exceptions)
 
     async def connect_real(self, device: Device, timeout: float, force_reconnect: bool):
-        """Used during ``Device.connect``.
+        """Use during [](#Device.connect) with `mock=False`.
 
-        This is called when a previous connect has not been done, or has been
-        done in a different mock more. It should connect the Device and all its
-        children.
+        This is called when there is no cached connect done in `mock=False`
+        mode. It connects the Device and all its children in real mode in parallel.
         """
         # Connect in parallel, gathering up NotConnected errors
         coros = {
@@ -61,11 +65,15 @@ class DeviceConnector:
 
 
 class Device(HasName):
-    """Common base class for all Ophyd Async Devices."""
+    """Common base class for all Ophyd Async Devices.
+
+    :param name: Optional name of the Device
+    :param connector: Optional DeviceConnector instance to use at connect()
+    """
 
-    _name: str = ""
-    #: The parent Device if it exists
     parent: Device | None = None
+    """The parent Device if it exists"""
+    _name: str = ""
     # None if connect hasn't started, a Task if it has
     _connect_task: asyncio.Task | None = None
     # The mock if we have connected in mock mode
@@ -83,7 +91,7 @@ class Device(HasName):
 
     @property
     def name(self) -> str:
-        """Return the name of the Device"""
+        """Return the name of the Device."""
         return self._name
 
     @cached_property
@@ -91,24 +99,26 @@ class Device(HasName):
         return {}
 
     def children(self) -> Iterator[tuple[str, Device]]:
+        """For each attribute that is a Device, yield the name and Device.
+
+        :yields: `(attr_name, attr)` for each child attribute that is a Device.
+        """
         yield from self._child_devices.items()
 
     @cached_property
     def log(self) -> LoggerAdapter:
+        """Return a logger configured with the device name."""
         return LoggerAdapter(
             getLogger("ophyd_async.devices"), {"ophyd_async_device_name": self.name}
         )
 
     def set_name(self, name: str, *, child_name_separator: str | None = None) -> None:
-        """Set ``self.name=name`` and each ``self.child.name=name+"-child"``.
-
-        Parameters
-        ----------
-        name:
-            New name to set
-        child_name_separator:
-            Use this as a separator instead of "-". Use "_" instead to make the same
-            names as the equivalent ophyd sync device.
+        """Set `self.name=name` and each `self.child.name=name+"-child"`.
+
+        :param name: New name to set.
+        :param child_name_separator:
+            Use this as a separator instead of "-". Use "_" instead to make the
+            same names as the equivalent ophyd sync device.
         """
         self._name = name
         if child_name_separator:
@@ -147,21 +157,26 @@ class Device(HasName):
         timeout: float = DEFAULT_TIMEOUT,
         force_reconnect: bool = False,
     ) -> None:
-        """Connect self and all child Devices.
-
-        Contains a timeout that gets propagated to child.connect methods.
-
-        Parameters
-        ----------
-        mock:
-            If True then use ``MockSignalBackend`` for all Signals
-        timeout:
-            Time to wait before failing with a TimeoutError.
+        """Connect the device and all child devices.
+
+        Successful connects will be cached so subsequent calls will return
+        immediately. Contains a timeout that gets propagated to child.connect
+        methods.
+
+        :param mock:
+            If True then use [](#MockSignalBackend) for all Signals. If passed a
+            [](#LazyMock) then pass this down for use within the Signals,
+            otherwise create one.
+        :param timeout: Time to wait before failing with a TimeoutError.
+        :param force_reconnect:
+            If True, force a reconnect even if the last connect succeeded.
         """
-        assert hasattr(self, "_connector"), (
-            f"{self}: doesn't have attribute `_connector`,"
-            " did you call `super().__init__` in your `__init__` method?"
-        )
+        if not hasattr(self, "_connector"):
+            msg = (
+                f"{self}: doesn't have attribute `_connector`,"
+                " did you call `super().__init__` in your `__init__` method?"
+            )
+            raise RuntimeError(msg)
         if mock:
             # Always connect in mock mode serially
             if isinstance(mock, LazyMock):
@@ -182,7 +197,9 @@ class Device(HasName):
                 self._mock = None
                 coro = self._connector.connect_real(self, timeout, force_reconnect)
                 self._connect_task = asyncio.create_task(coro)
-            assert self._connect_task, "Connect task not created, this shouldn't happen"
+            if not self._connect_task:
+                msg = "Connect task not created, this shouldn't happen"
+                raise RuntimeError(msg)
             # Wait for it to complete
             await self._connect_task
 
@@ -201,12 +218,9 @@ DeviceT = TypeVar("DeviceT", bound=Device)
 
 
 class DeviceVector(MutableMapping[int, DeviceT], Device):
-    """
-    Defines device components with indices.
+    """Defines a dictionary of Device children with arbitrary integer keys.
 
-    In the below example, foos becomes a dictionary on the parent device
-    at runtime, so parent.foos[2] returns a FooDevice. For example usage see
-    :class:`~ophyd_async.epics.demo.DynamicSensorGroup`
+    :see-also: [](#implementing-devices) for examples of how to use this class.
     """
 
     def __init__(
@@ -232,8 +246,12 @@ class DeviceVector(MutableMapping[int, DeviceT], Device):
     def __setitem__(self, key: int, value: DeviceT) -> None:
         # Check the types on entry to dict to make sure we can't accidentally
         # make a non-integer named child
-        assert isinstance(key, int), f"Expected int, got {key}"
-        assert isinstance(value, Device), f"Expected Device, got {value}"
+        if not isinstance(key, int):
+            msg = f"Expected int, got {key}"
+            raise TypeError(msg)
+        if not isinstance(value, Device):
+            msg = f"Expected Device, got {value}"
+            raise TypeError(msg)
         self._children[key] = value
         value.parent = self
 
@@ -254,94 +272,49 @@ class DeviceVector(MutableMapping[int, DeviceT], Device):
         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 DeviceProcessor:
+    """Sync/Async Context Manager that finds all the Devices declared within it.
 
+    Used in `init_devices`
     """
 
-    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):
-        """Walk up until we find a stack frame that doesn't have us as 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
         except ValueError:
             _, _, tb = sys.exc_info()
-            assert tb, "Can't get traceback, this shouldn't happen"
+            if not tb:
+                msg = "Can't get traceback, this shouldn't happen"
+                raise RuntimeError(msg)  # noqa: B904
             caller_frame = tb.tb_frame
             while caller_frame.f_locals.get("self", None) is self:
                 caller_frame = caller_frame.f_back
-                assert (
-                    caller_frame
-                ), "No previous frame to the one with self in it, this shouldn't happen"
-            return caller_frame.f_locals
+                if not caller_frame:
+                    msg = (
+                        "No previous frame to the one with self in it, "
+                        "this shouldn't happen"
+                    )
+                    raise RuntimeError(  # noqa: B904
+                        msg
+                    )
+            return caller_frame.f_locals.copy()
 
-    def __enter__(self) -> DeviceCollector:
+    def __enter__(self) -> DeviceProcessor:
         # 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) -> DeviceProcessor:
         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 +323,7 @@ class DeviceCollector:
                 "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 +333,58 @@ class DeviceCollector:
                 "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 init_devices(
+    set_name=True,
+    child_name_separator: str = "-",
+    connect=True,
+    mock=False,
+    timeout: float = 10.0,
+):
+    """Auto initialize top level Device instances: to be used as a context manager.
+
+    :param set_name:
+        If True, call `device.set_name(variable_name)` on all Devices created
+        within the context manager that have an empty `name`.
+    :param child_name_separator: Separator for child names if `set_name` is True.
+    :param connect:
+        If True, call `device.connect(mock, timeout)` in parallel on all Devices
+        created within the context manager.
+    :param mock: If True, connect Signals in mock mode.
+    :param timeout: How long to wait for connect before logging an exception.
+    :raises RuntimeError: If used inside a plan, use [](#ensure_connected) instead.
+    :raises NotConnected: If devices could not be connected.
+
+    For example, to connect and name 2 motors in parallel:
+    ```python
+    [async] with 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 DeviceProcessor(process_devices)

ophyd_async/core/_device_filler.py

@@ -16,7 +16,7 @@ from typing import (
 )
 
 from ._device import Device, DeviceConnector, DeviceVector
-from ._signal import Signal, SignalX
+from ._signal import Ignore, Signal, SignalX
 from ._signal_backend import SignalBackend, SignalDatatype
 from ._utils import get_origin_class
 
@@ -33,12 +33,20 @@ def _get_datatype(annotation: Any) -> type | None:
     args = get_args(annotation)
     if len(args) == 1 and get_origin_class(args[0]):
         return args[0]
+    return None
 
 
 def _logical(name: UniqueName) -> LogicalName:
     return LogicalName(name.rstrip("_"))
 
 
+def _check_device_annotation(annotation: Any) -> DeviceAnnotation:
+    if not isinstance(annotation, DeviceAnnotation):
+        msg = f"Annotation {annotation} is not a DeviceAnnotation"
+        raise TypeError(msg)
+    return annotation
+
+
 @runtime_checkable
 class DeviceAnnotation(Protocol):
     @abstractmethod
@@ -46,6 +54,13 @@ class DeviceAnnotation(Protocol):
 
 
 class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
+    """For filling signals on introspected devices.
+
+    :param device: The device to fill.
+    :param signal_backend_factory: A callable that returns a SignalBackend.
+    :param device_connector_factory: A callable that returns a DeviceConnector.
+    """
+
     def __init__(
         self,
         device: Device,
@@ -61,6 +76,7 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         self._extras: dict[UniqueName, Sequence[Any]] = {}
         self._signal_datatype: dict[LogicalName, type | None] = {}
         self._vector_device_type: dict[LogicalName, type[Device] | None] = {}
+        self.ignored_signals: set[str] = set()
         # Backends and Connectors stored ready for the connection phase
         self._unfilled_backends: dict[
             LogicalName, tuple[SignalBackendT, type[Signal]]
@@ -101,6 +117,8 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         # Get hints with Annotated for wrapping signals and backends
         extra_hints = get_type_hints(cls, include_extras=True)
         for attr_name, annotation in hints.items():
+            if annotation is Ignore:
+                self.ignored_signals.add(attr_name)
             name = UniqueName(attr_name)
             origin = get_origin_class(annotation)
             if (
@@ -131,6 +149,7 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
                 self._uncreated_devices[name] = origin
 
     def check_created(self):
+        """Check that all Signals and Devices declared in annotations are created."""
         uncreated = sorted(set(self._uncreated_signals).union(self._uncreated_devices))
         if uncreated:
             raise RuntimeError(
@@ -141,6 +160,21 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         self,
         filled=True,
     ) -> Iterator[tuple[SignalBackendT, list[Any]]]:
+        """Create all Signals from annotations.
+
+        :param filled:
+            If True then the Signals created should be considered already filled
+            with connection data. If False then `fill_child_signal` needs
+            calling at device connection time before the signal can be
+            connected.
+        :yields: `(backend, extras)`
+            The `SignalBackend` that has been created for this Signal, and the
+            list of extra annotations that could be used to customize it. For
+            example an `EpicsDeviceConnector` consumes `PvSuffix` extras to set the
+            write_pv of the backend. Any unhandled extras should be left on the
+            list so this class can handle them, e.g. `StandardReadableFormat`
+            instances.
+        """
         for name in list(self._uncreated_signals):
             child_type = self._uncreated_signals.pop(name)
             backend = self._signal_backend_factory(
@@ -150,8 +184,8 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
             yield backend, extras
             signal = child_type(backend)
             for anno in extras:
-                assert isinstance(anno, DeviceAnnotation), anno
-                anno(self._device, signal)
+                device_annotation = _check_device_annotation(annotation=anno)
+                device_annotation(self._device, signal)
             setattr(self._device, name, signal)
             dest = self._filled_backends if filled else self._unfilled_backends
             dest[_logical(name)] = (backend, child_type)
@@ -160,6 +194,17 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         self,
         filled=True,
     ) -> Iterator[tuple[DeviceConnectorT, list[Any]]]:
+        """Create all Signals from annotations.
+
+        :param filled:
+            If True then the Devices created should be considered already filled
+            with connection data. If False then `fill_child_device` needs
+            calling at parent device connection time before the child Device can
+            be connected.
+        :yields: `(connector, extras)`
+            The `DeviceConnector` that has been created for this Signal, and the list of
+            extra annotations that could be used to customize it.
+        """
         for name in list(self._uncreated_devices):
             child_type = self._uncreated_devices.pop(name)
             connector = self._device_connector_factory()
@@ -167,15 +212,21 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
             yield connector, extras
             device = child_type(connector=connector)
             for anno in extras:
-                assert isinstance(anno, DeviceAnnotation), anno
-                anno(self._device, device)
+                device_annotation = _check_device_annotation(annotation=anno)
+                device_annotation(self._device, device)
             setattr(self._device, name, device)
             dest = self._filled_connectors if filled else self._unfilled_connectors
             dest[_logical(name)] = connector
 
     def create_device_vector_entries_to_mock(self, num: int):
+        """Create num entries for each `DeviceVector`.
+
+        This is used when the Device is being connected in mock mode.
+        """
         for name, cls in self._vector_device_type.items():
-            assert cls, "Shouldn't happen"
+            if not cls:
+                msg = "Malformed device vector"
+                raise TypeError(msg)
             for i in range(1, num + 1):
                 if issubclass(cls, Signal):
                     self.fill_child_signal(name, cls, i)
@@ -185,6 +236,11 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
                     self._raise(name, f"Can't make {cls}")
 
     def check_filled(self, source: str):
+        """Check that all the created Signals and Devices are filled.
+
+        :param source: The source of the data that should have done the filling, for
+                       reporting as an error message
+        """
         unfilled = sorted(set(self._unfilled_connectors).union(self._unfilled_backends))
         if unfilled:
             raise RuntimeError(
@@ -207,6 +263,15 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         signal_type: type[Signal],
         vector_index: int | None = None,
     ) -> SignalBackendT:
+        """Mark a Signal as filled, and return its backend for filling.
+
+        :param name:
+            The name without trailing underscore, the name in the control system
+        :param signal_type:
+            One of the types `SignalR`, `SignalW`, `SignalRW` or `SignalX`
+        :param vector_index: If the child is in a `DeviceVector` then what index is it
+        :return: The SignalBackend for the filled Signal.
+        """
         name = cast(LogicalName, name)
         if name in self._unfilled_backends:
             # We made it above
@@ -242,6 +307,14 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         device_type: type[Device] = Device,
         vector_index: int | None = None,
     ) -> DeviceConnectorT:
+        """Mark a Device as filled, and return its connector for filling.
+
+        :param name:
+            The name without trailing underscore, the name in the control system
+        :param device_type: The `Device` subclass to be created
+        :param vector_index: If the child is in a `DeviceVector` then what index is it
+        :return: The DeviceConnector for the filled Device.
+        """
         name = cast(LogicalName, name)
         if name in self._unfilled_connectors:
             # We made it above
@@ -254,9 +327,9 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
             # We need to add a new entry to a DeviceVector
             vector = self._ensure_device_vector(name)
             vector_device_type = self._vector_device_type[name] or device_type
-            assert issubclass(
-                vector_device_type, Device
-            ), f"{vector_device_type} is not a Device"
+            if not issubclass(vector_device_type, Device):
+                msg = f"{vector_device_type} is not a Device"
+                raise TypeError(msg)
             connector = self._device_connector_factory()
             vector[vector_index] = vector_device_type(connector=connector)
         elif child := getattr(self._device, name, None):

ophyd_async/core/_flyer.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Generic
+from typing import Any, Generic
 
 from bluesky.protocols import Flyable, Preparable, Stageable
 
@@ -9,21 +9,26 @@ from ._utils import T
 
 
 class FlyerController(ABC, Generic[T]):
+    """Base class for controlling 'flyable' devices.
+
+    [`bluesky.protocols.Flyable`](#bluesky.protocols.Flyable).
+    """
+
     @abstractmethod
-    async def prepare(self, value: T):
-        """Move to the start of the flyscan"""
+    async def prepare(self, value: T) -> Any:
+        """Move to the start of the flyscan."""
 
     @abstractmethod
     async def kickoff(self):
-        """Start the flyscan"""
+        """Start the flyscan."""
 
     @abstractmethod
     async def complete(self):
-        """Block until the flyscan is done"""
+        """Block until the flyscan is done."""
 
     @abstractmethod
     async def stop(self):
-        """Stop flying and wait everything to be stopped"""
+        """Stop flying and wait everything to be stopped."""
 
 
 class StandardFlyer(
@@ -33,6 +38,11 @@ class StandardFlyer(
     Flyable,
     Generic[T],
 ):
+    """Base class for 'flyable' devices.
+
+    [`bluesky.protocols.Flyable`](#bluesky.protocols.Flyable).
+    """
+
     def __init__(
         self,
         trigger_logic: FlyerController[T],
@@ -54,7 +64,6 @@ class StandardFlyer(
         await self._trigger_logic.stop()
 
     def prepare(self, value: T) -> AsyncStatus:
-        """Setup trajectories"""
         return AsyncStatus(self._prepare(value))
 
     async def _prepare(self, value: T) -> None: