ophyd-async 0.7.0__py3-none-any.whl → 0.8.0a2__py3-none-any.whl

ophyd_async/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.0'
-__version_tuple__ = version_tuple = (0, 7, 0)
+__version__ = version = '0.8.0a2'
+__version_tuple__ = version_tuple = (0, 8, 0)

ophyd_async/core/__init__.py

@@ -5,7 +5,8 @@ from ._detector import (
     StandardDetector,
     TriggerInfo,
 )
-from ._device import Device, DeviceCollector, DeviceVector
+from ._device import Device, DeviceCollector, DeviceConnector, DeviceVector
+from ._device_filler import DeviceFiller
 from ._device_save_loader import (
     all_at_once,
     get_signal_values,
@@ -22,6 +23,7 @@ from ._log import config_ophyd_async_logging
 from ._mock_signal_backend import MockSignalBackend
 from ._mock_signal_utils import (
     callback_on_mock_put,
+    get_mock,
     get_mock_put,
     mock_puts_blocked,
     reset_mock_put_calls,
@@ -62,9 +64,11 @@ from ._signal import (
     wait_for_value,
 )
 from ._signal_backend import (
-    RuntimeSubsetEnum,
+    Array1D,
     SignalBackend,
-    SubsetEnum,
+    SignalDatatype,
+    SignalDatatypeT,
+    make_datakey,
 )
 from ._soft_signal_backend import SignalMetadata, SoftSignalBackend
 from ._status import AsyncStatus, WatchableAsyncStatus, completed_status
@@ -73,14 +77,17 @@ from ._utils import (
     CALCULATE_TIMEOUT,
     DEFAULT_TIMEOUT,
     CalculatableTimeout,
+    Callback,
     NotConnected,
-    ReadingValueCallback,
+    Reference,
+    StrictEnum,
+    SubsetEnum,
     T,
     WatcherUpdate,
     get_dtype,
+    get_enum_cls,
     get_unique,
     in_micros,
-    is_pydantic_model,
     wait_for_connection,
 )
 
@@ -91,8 +98,10 @@ __all__ = [
     "StandardDetector",
     "TriggerInfo",
     "Device",
+    "DeviceConnector",
     "DeviceCollector",
     "DeviceVector",
+    "DeviceFiller",
     "all_at_once",
     "get_signal_values",
     "load_device",
@@ -108,6 +117,7 @@ __all__ = [
     "config_ophyd_async_logging",
     "MockSignalBackend",
     "callback_on_mock_put",
+    "get_mock",
     "get_mock_put",
     "mock_puts_blocked",
     "reset_mock_put_calls",
@@ -146,25 +156,30 @@ __all__ = [
     "soft_signal_r_and_setter",
     "soft_signal_rw",
     "wait_for_value",
-    "RuntimeSubsetEnum",
+    "Array1D",
     "SignalBackend",
+    "make_datakey",
+    "StrictEnum",
     "SubsetEnum",
+    "SignalDatatype",
+    "SignalDatatypeT",
     "SignalMetadata",
     "SoftSignalBackend",
     "AsyncStatus",
     "WatchableAsyncStatus",
     "DEFAULT_TIMEOUT",
     "CalculatableTimeout",
+    "Callback",
     "CALCULATE_TIMEOUT",
     "NotConnected",
-    "ReadingValueCallback",
+    "Reference",
     "Table",
     "T",
     "WatcherUpdate",
     "get_dtype",
+    "get_enum_cls",
     "get_unique",
     "in_micros",
-    "is_pydantic_model",
     "wait_for_connection",
     "completed_status",
 ]

ophyd_async/core/_detector.py

@@ -4,11 +4,7 @@ import asyncio
 import time
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Sequence
-from enum import Enum
 from functools import cached_property
-from typing import (
-    Generic,
-)
 
 from bluesky.protocols import (
     Collectable,
@@ -23,14 +19,14 @@ from bluesky.protocols import (
 from event_model import DataKey
 from pydantic import BaseModel, Field, NonNegativeInt, computed_field
 
-from ._device import Device
+from ._device import Device, DeviceConnector
 from ._protocol import AsyncConfigurable, AsyncReadable
 from ._signal import SignalR
 from ._status import AsyncStatus, WatchableAsyncStatus
-from ._utils import DEFAULT_TIMEOUT, T, WatcherUpdate, merge_gathered_dicts
+from ._utils import DEFAULT_TIMEOUT, StrictEnum, WatcherUpdate, merge_gathered_dicts
 
 
-class DetectorTrigger(str, Enum):
+class DetectorTrigger(StrictEnum):
     """Type of mechanism for triggering a detector to take frames"""
 
     #: Detector generates internal trigger for given rate
@@ -172,7 +168,6 @@ class StandardDetector(
     Flyable,
     Collectable,
     WritesStreamAssets,
-    Generic[T],
 ):
     """
     Useful detector base class for step and fly scanning detectors.
@@ -185,6 +180,7 @@ class StandardDetector(
         writer: DetectorWriter,
         config_sigs: Sequence[SignalR] = (),
         name: str = "",
+        connector: DeviceConnector | None = None,
     ) -> None:
         """
         Constructor
@@ -213,8 +209,7 @@ class StandardDetector(
         self._completable_frames: int = 0
         self._number_of_triggers_iter: Iterator[int] | None = None
         self._initial_frame: int = 0
-
-        super().__init__(name)
+        super().__init__(name, connector=connector)
 
     @property
     def controller(self) -> DetectorController:

ophyd_async/core/_device.py

@@ -1,39 +1,80 @@
-"""Base device"""
+from __future__ import annotations
 
 import asyncio
 import sys
-from collections.abc import Coroutine, Generator, Iterator
-from functools import cached_property
+from collections.abc import Coroutine, Iterator, Mapping, MutableMapping
 from logging import LoggerAdapter, getLogger
-from typing import (
-    Any,
-    Optional,
-    TypeVar,
-)
+from typing import Any, TypeVar
+from unittest.mock import Mock
 
 from bluesky.protocols import HasName
 from bluesky.run_engine import call_in_bluesky_event_loop, in_bluesky_event_loop
 
+from ._protocol import Connectable
 from ._utils import DEFAULT_TIMEOUT, NotConnected, wait_for_connection
 
+_device_mocks: dict[Device, Mock] = {}
 
-class Device(HasName):
-    """Common base class for all Ophyd Async Devices.
 
-    By default, names and connects all Device children.
-    """
+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.
+
+        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
+        hints to the Device like so::
+
+            my_signal: SignalRW[int]
+            my_device: MyDevice
+
+        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__``.
+        """
+
+    async def connect(
+        self,
+        device: Device,
+        mock: bool | Mock,
+        timeout: float,
+        force_reconnect: bool,
+    ):
+        """Used during ``Device.connect``.
+
+        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.
+        """
+        coros = {}
+        for name, child_device in device.children():
+            child_mock = getattr(mock, name) if mock else mock  # Mock() or False
+            coros[name] = child_device.connect(
+                mock=child_mock, timeout=timeout, force_reconnect=force_reconnect
+            )
+        await wait_for_connection(**coros)
+
+
+class Device(HasName, Connectable):
+    """Common base class for all Ophyd Async Devices."""
 
     _name: str = ""
     #: The parent Device if it exists
-    parent: Optional["Device"] = None
+    parent: Device | None = None
     # None if connect hasn't started, a Task if it has
     _connect_task: asyncio.Task | None = None
+    # If not None, then this is the mock arg of the previous connect
+    # to let us know if we can reuse an existing connection
+    _connect_mock_arg: bool | None = None
 
-    # Used to check if the previous connect was mocked,
-    # if the next mock value differs then we fail
-    _previous_connect_was_mock = None
-
-    def __init__(self, name: str = "") -> None:
+    def __init__(
+        self, name: str = "", connector: DeviceConnector | None = None
+    ) -> None:
+        self._connector = connector or DeviceConnector()
         self.set_name(name)
 
     @property
@@ -41,13 +82,7 @@ class Device(HasName):
         """Return the name of the Device"""
         return self._name
 
-    @cached_property
-    def log(self):
-        return LoggerAdapter(
-            getLogger("ophyd_async.devices"), {"ophyd_async_device_name": self.name}
-        )
-
-    def children(self) -> Iterator[tuple[str, "Device"]]:
+    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
@@ -60,23 +95,32 @@ class Device(HasName):
         name:
             New name to set
         """
-
-        # Ensure self.log is recreated after a name change
-        if hasattr(self, "log"):
-            del self.log
-
         self._name = name
-        for attr_name, child in self.children():
-            child_name = f"{name}-{attr_name.rstrip('_')}" if name else ""
+        # Ensure self.log is recreated after a name change
+        self.log = LoggerAdapter(
+            getLogger("ophyd_async.devices"), {"ophyd_async_device_name": self.name}
+        )
+        for child_name, child in self.children():
+            child_name = f"{self.name}-{child_name.strip('_')}" if self.name else ""
             child.set_name(child_name)
-            child.parent = self
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name == "parent":
+            if self.parent not in (value, None):
+                raise TypeError(
+                    f"Cannot set the parent of {self} to be {value}: "
+                    f"it is already a child of {self.parent}"
+                )
+        elif isinstance(value, Device):
+            value.parent = self
+        return super().__setattr__(name, value)
 
     async def connect(
         self,
-        mock: bool = False,
+        mock: bool | Mock = False,
         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.
@@ -88,41 +132,32 @@ class Device(HasName):
         timeout:
             Time to wait before failing with a TimeoutError.
         """
-
-        if (
-            self._previous_connect_was_mock is not None
-            and self._previous_connect_was_mock != mock
-        ):
-            raise RuntimeError(
-                f"`connect(mock={mock})` called on a `Device` where the previous "
-                f"connect was `mock={self._previous_connect_was_mock}`. Changing mock "
-                "value between connects is not permitted."
-            )
-        self._previous_connect_was_mock = mock
-
-        # If previous connect with same args has started and not errored, can use it
-        can_use_previous_connect = self._connect_task and not (
-            self._connect_task.done() and self._connect_task.exception()
+        uses_mock = bool(mock)
+        can_use_previous_connect = (
+            uses_mock is self._connect_mock_arg
+            and self._connect_task
+            and not (self._connect_task.done() and self._connect_task.exception())
         )
+        if mock is True:
+            mock = Mock()  # create a new Mock if one not provided
         if force_reconnect or not can_use_previous_connect:
-            # Kick off a connection
-            coros = {
-                name: child_device.connect(
-                    mock, timeout=timeout, force_reconnect=force_reconnect
-                )
-                for name, child_device in self.children()
-            }
-            self._connect_task = asyncio.create_task(wait_for_connection(**coros))
+            self._connect_mock_arg = uses_mock
+            if self._connect_mock_arg:
+                _device_mocks[self] = mock
+            coro = self._connector.connect(
+                device=self, mock=mock, timeout=timeout, force_reconnect=force_reconnect
+            )
+            self._connect_task = asyncio.create_task(coro)
 
         assert self._connect_task, "Connect task not created, this shouldn't happen"
         # Wait for it to complete
         await self._connect_task
 
 
-VT = TypeVar("VT", bound=Device)
+DeviceT = TypeVar("DeviceT", bound=Device)
 
 
-class DeviceVector(dict[int, VT], Device):
+class DeviceVector(MutableMapping[int, DeviceT], Device):
     """
     Defines device components with indices.
 
@@ -131,10 +166,48 @@ class DeviceVector(dict[int, VT], Device):
     :class:`~ophyd_async.epics.demo.DynamicSensorGroup`
     """
 
-    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
+    def __init__(
+        self,
+        children: Mapping[int, DeviceT],
+        name: str = "",
+    ) -> None:
+        self._children = dict(children)
+        super().__init__(name=name)
+
+    def __setattr__(self, name: str, child: Any) -> None:
+        if name != "parent" and isinstance(child, Device):
+            raise AttributeError(
+                "DeviceVector can only have integer named children, "
+                "set via device_vector[i] = child"
+            )
+        super().__setattr__(name, child)
+
+    def __getitem__(self, key: int) -> DeviceT:
+        return self._children[key]
+
+    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}"
+        self._children[key] = value
+        value.parent = self
+
+    def __delitem__(self, key: int) -> None:
+        del self._children[key]
+
+    def __iter__(self) -> Iterator[int]:
+        yield from self._children
+
+    def __len__(self) -> int:
+        return len(self._children)
+
+    def children(self) -> Iterator[tuple[str, Device]]:
+        for key, child in self._children.items():
+            yield str(key), child
+
+    def __hash__(self):  # to allow DeviceVector to be used as dict keys and in sets
+        return hash(id(self))
 
 
 class DeviceCollector:
@@ -195,12 +268,12 @@ class DeviceCollector:
                 ), "No previous frame to the one with self in it, this shouldn't happen"
             return caller_frame.f_locals
 
-    def __enter__(self) -> "DeviceCollector":
+    def __enter__(self) -> DeviceCollector:
         # Stash the names that were defined before we were called
         self._names_on_enter = set(self._caller_locals())
         return self
 
-    async def __aenter__(self) -> "DeviceCollector":
+    async def __aenter__(self) -> DeviceCollector:
         return self.__enter__()
 
     async def _on_exit(self) -> None:

ophyd_async/core/_device_filler.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from typing import (
+    Generic,
+    NoReturn,
+    TypeVar,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from ._device import Device, DeviceConnector, DeviceVector
+from ._signal import Signal, SignalX
+from ._signal_backend import SignalBackend, SignalDatatype
+from ._utils import get_origin_class
+
+
+def _strip_number_from_string(string: str) -> tuple[str, int | None]:
+    match = re.match(r"(.*?)(\d*)$", string)
+    assert match
+
+    name = match.group(1)
+    number = match.group(2) or None
+    if number is None:
+        return name, None
+    else:
+        return name, int(number)
+
+
+SignalBackendT = TypeVar("SignalBackendT", bound=SignalBackend)
+DeviceConnectorT = TypeVar("DeviceConnectorT", bound=DeviceConnector)
+
+
+class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
+    def __init__(
+        self,
+        device: Device,
+        signal_backend_factory: Callable[[type[SignalDatatype] | None], SignalBackendT],
+        device_connector_factory: Callable[[], DeviceConnectorT],
+    ):
+        self._device = device
+        self._signal_backend_factory = signal_backend_factory
+        self._device_connector_factory = device_connector_factory
+        self._vectors: dict[str, DeviceVector] = {}
+        self._vector_device_type: dict[str, type[Device] | None] = {}
+        self._signal_backends: dict[str, tuple[SignalBackendT, type[Signal]]] = {}
+        self._device_connectors: dict[str, DeviceConnectorT] = {}
+        # Get type hints on the class, not the instance
+        # https://github.com/python/cpython/issues/124840
+        self._annotations = get_type_hints(type(device))
+        for name, annotation in self._annotations.items():
+            # names have a trailing underscore if the clash with a bluesky verb,
+            # so strip this off to get what the CS will provide
+            stripped_name = name.rstrip("_")
+            origin = get_origin_class(annotation)
+            if name == "parent" or name.startswith("_") or not origin:
+                # Ignore
+                pass
+            elif issubclass(origin, Signal):
+                # SignalX doesn't need datatype, all others need one
+                datatype = self.get_datatype(name)
+                if origin != SignalX and datatype is None:
+                    self._raise(
+                        name,
+                        f"Expected SignalX or SignalR/W/RW[type], got {annotation}",
+                    )
+                self._signal_backends[stripped_name] = (
+                    self.make_child_signal(name, origin),
+                    origin,
+                )
+            elif origin == DeviceVector:
+                # DeviceVector needs a type of device
+                args = get_args(annotation) or [None]
+                child_origin = get_origin(args[0]) or args[0]
+                if child_origin is None or not issubclass(child_origin, Device):
+                    self._raise(
+                        name,
+                        f"Expected DeviceVector[SomeDevice], got {annotation}",
+                    )
+                self.make_device_vector(name, child_origin)
+            elif issubclass(origin, Device):
+                self._device_connectors[stripped_name] = self.make_child_device(
+                    name, origin
+                )
+
+    def unfilled(self) -> set[str]:
+        return set(self._device_connectors).union(self._signal_backends)
+
+    def _raise(self, name: str, error: str) -> NoReturn:
+        raise TypeError(f"{type(self._device).__name__}.{name}: {error}")
+
+    def make_device_vector(self, name: str, device_type: type[Device] | None):
+        self._vectors[name] = DeviceVector({})
+        self._vector_device_type[name] = device_type
+        setattr(self._device, name, self._vectors[name])
+
+    def make_device_vectors(self, names: list[str]):
+        basenames: dict[str, set[int]] = {}
+        for name in names:
+            basename, number = _strip_number_from_string(name)
+            if number is not None:
+                basenames.setdefault(basename, set()).add(number)
+        for basename, numbers in basenames.items():
+            # If contiguous numbers starting at 1 then it's a device vector
+            length = len(numbers)
+            if length > 1 and numbers == set(range(1, length + 1)):
+                # DeviceVector needs a type of device
+                self.make_device_vector(basename, None)
+
+    def get_datatype(self, name: str) -> type[SignalDatatype] | None:
+        # Get dtype from SignalRW[dtype] or DeviceVector[SignalRW[dtype]]
+        basename, _ = _strip_number_from_string(name)
+        if basename in self._vectors:
+            # We decided to put it in a device vector, so get datatype from that
+            annotation = self._annotations.get(basename, None)
+            if annotation:
+                annotation = get_args(annotation)[0]
+        else:
+            # It's not a device vector, so get it from the full name
+            annotation = self._annotations.get(name, None)
+        args = get_args(annotation)
+        if args and get_origin_class(args[0]):
+            return args[0]
+
+    def make_child_signal(self, name: str, signal_type: type[Signal]) -> SignalBackendT:
+        if name in self._signal_backends:
+            # We made it above
+            backend, expected_signal_type = self._signal_backends.pop(name)
+        else:
+            # We need to make a new one
+            basename, number = _strip_number_from_string(name)
+            child = getattr(self._device, name, None)
+            backend = self._signal_backend_factory(self.get_datatype(name))
+            signal = signal_type(backend)
+            if basename in self._vectors and isinstance(number, int):
+                # We need to add a new entry to an existing DeviceVector
+                expected_signal_type = self._vector_device_type[basename] or signal_type
+                self._vectors[basename][number] = signal
+            elif child is None:
+                # We need to add a new child to the top level Device
+                expected_signal_type = signal_type
+                setattr(self._device, name, signal)
+            else:
+                self._raise(name, f"Cannot make child as it would shadow {child}")
+        if signal_type is not expected_signal_type:
+            self._raise(
+                name,
+                f"is a {signal_type.__name__} not a {expected_signal_type.__name__}",
+            )
+        return backend
+
+    def make_child_device(
+        self, name: str, device_type: type[Device] = Device
+    ) -> DeviceConnectorT:
+        basename, number = _strip_number_from_string(name)
+        child = getattr(self._device, name, None)
+        if connector := self._device_connectors.pop(name, None):
+            # We made it above
+            return connector
+        elif basename in self._vectors and isinstance(number, int):
+            # We need to add a new entry to an existing DeviceVector
+            vector_device_type = self._vector_device_type[basename] or device_type
+            assert issubclass(
+                vector_device_type, Device
+            ), f"{vector_device_type} is not a Device"
+            connector = self._device_connector_factory()
+            device = vector_device_type(connector=connector)
+            self._vectors[basename][number] = device
+        elif child is None:
+            # We need to add a new child to the top level Device
+            connector = self._device_connector_factory()
+            device = device_type(connector=connector)
+            setattr(self._device, name, device)
+        else:
+            self._raise(name, f"Cannot make child as it would shadow {child}")
+        connector.create_children_from_annotations(device)
+        return connector
+
+    def make_soft_device_vector_entries(self, num: int):
+        for basename, cls in self._vector_device_type.items():
+            assert cls, "Shouldn't happen"
+            for i in range(num):
+                name = f"{basename}{i + 1}"
+                if issubclass(cls, Signal):
+                    self.make_child_signal(name, cls)
+                elif issubclass(cls, Device):
+                    self.make_child_device(name, cls)
+                else:
+                    self._raise(name, f"Can't make {cls}")

ophyd_async/core/_device_save_loader.py

@@ -27,11 +27,8 @@ def pydantic_model_abstraction_representer(
     return dumper.represent_data(model.model_dump(mode="python"))
 
 
-class OphydDumper(yaml.Dumper):
-    def represent_data(self, data: Any) -> Any:
-        if isinstance(data, Enum):
-            return self.represent_data(data.value)
-        return super().represent_data(data)
+def enum_representer(dumper: yaml.Dumper, enum: Enum) -> yaml.Node:
+    return dumper.represent_data(enum.value)
 
 
 def get_signal_values(
@@ -74,7 +71,7 @@ def get_signal_values(
         for key, value in zip(selected_signals, selected_values, strict=False)
     }
     # Ignored values place in with value None so we know which ones were ignored
-    named_values.update({key: None for key in ignore})
+    named_values.update(dict.fromkeys(ignore))
     return named_values
 
 
@@ -111,6 +108,7 @@ def walk_rw_signals(
         path_prefix = ""
 
     signals: dict[str, SignalRW[Any]] = {}
+
     for attr_name, attr in device.children():
         dot_path = f"{path_prefix}{attr_name}"
         if type(attr) is SignalRW:
@@ -145,9 +143,10 @@ def save_to_yaml(phases: Sequence[dict[str, Any]], save_path: str | Path) -> Non
         pydantic_model_abstraction_representer,
         Dumper=yaml.Dumper,
     )
+    yaml.add_multi_representer(Enum, enum_representer, Dumper=yaml.Dumper)
 
     with open(save_path, "w") as file:
-        yaml.dump(phases, file, Dumper=OphydDumper, default_flow_style=False)
+        yaml.dump(phases, file)
 
 
 def load_from_yaml(save_path: str) -> Sequence[dict[str, Any]]: