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

ophyd_async/core/_status.py

@@ -13,6 +13,7 @@ from typing import (
 
 from bluesky.protocols import Status
 
+from ._device import Device
 from ._protocol import Watcher
 from ._utils import Callback, P, T, WatcherUpdate
 
@@ -23,13 +24,14 @@ WAS = TypeVar("WAS", bound="WatchableAsyncStatus")
 class AsyncStatusBase(Status):
     """Convert asyncio awaitable to bluesky Status interface"""
 
-    def __init__(self, awaitable: Coroutine | asyncio.Task):
+    def __init__(self, awaitable: Coroutine | asyncio.Task, name: str | None = None):
         if isinstance(awaitable, asyncio.Task):
             self.task = awaitable
         else:
             self.task = asyncio.create_task(awaitable)
         self.task.add_done_callback(self._run_callbacks)
         self._callbacks: list[Callback[Status]] = []
+        self._name = name
 
     def __await__(self):
         return self.task.__await__()
@@ -76,7 +78,11 @@ class AsyncStatusBase(Status):
                 status = "done"
         else:
             status = "pending"
-        return f"<{type(self).__name__}, task: {self.task.get_coro()}, {status}>"
+        device_str = f"device: {self._name}, " if self._name else ""
+        return (
+            f"<{type(self).__name__}, {device_str}"
+            f"task: {self.task.get_coro()}, {status}>"
+        )
 
     __str__ = __repr__
 
@@ -90,7 +96,11 @@ class AsyncStatus(AsyncStatusBase):
 
         @functools.wraps(f)
         def wrap_f(*args: P.args, **kwargs: P.kwargs) -> AS:
-            return cls(f(*args, **kwargs))
+            if args and isinstance(args[0], Device):
+                name = args[0].name
+            else:
+                name = None
+            return cls(f(*args, **kwargs), name=name)
 
         # type is actually functools._Wrapped[P, Awaitable, P, AS]
         # but functools._Wrapped is not necessarily available
@@ -100,11 +110,13 @@ class AsyncStatus(AsyncStatusBase):
 class WatchableAsyncStatus(AsyncStatusBase, Generic[T]):
     """Convert AsyncIterator of WatcherUpdates to bluesky Status interface."""
 
-    def __init__(self, iterator: AsyncIterator[WatcherUpdate[T]]):
+    def __init__(
+        self, iterator: AsyncIterator[WatcherUpdate[T]], name: str | None = None
+    ):
         self._watchers: list[Watcher] = []
         self._start = time.monotonic()
         self._last_update: WatcherUpdate[T] | None = None
-        super().__init__(self._notify_watchers_from(iterator))
+        super().__init__(self._notify_watchers_from(iterator), name)
 
     async def _notify_watchers_from(self, iterator: AsyncIterator[WatcherUpdate[T]]):
         async for update in iterator:
@@ -136,7 +148,11 @@ class WatchableAsyncStatus(AsyncStatusBase, Generic[T]):
 
         @functools.wraps(f)
         def wrap_f(*args: P.args, **kwargs: P.kwargs) -> WAS:
-            return cls(f(*args, **kwargs))
+            if args and isinstance(args[0], Device):
+                name = args[0].name
+            else:
+                name = None
+            return cls(f(*args, **kwargs), name=name)
 
         return cast(Callable[P, WAS], wrap_f)
 

ophyd_async/core/_table.py

@@ -1,8 +1,13 @@
-from enum import Enum
-from typing import TypeVar, get_args, get_origin
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Annotated, Any, TypeVar, get_origin
 
 import numpy as np
-from pydantic import BaseModel, ConfigDict, model_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic_numpy.helper.annotation import NpArrayPydanticAnnotation
+
+from ._utils import get_dtype
 
 TableSubclass = TypeVar("TableSubclass", bound="Table")
 
@@ -14,37 +19,46 @@ def _concat(value1, value2):
         return value1 + value2
 
 
+def _make_default_factory(dtype: np.dtype) -> Callable[[], np.ndarray]:
+    def numpy_array_default_factory() -> np.ndarray:
+        return np.array([], dtype)
+
+    return numpy_array_default_factory
+
+
 class Table(BaseModel):
     """An abstraction of a Table of str to numpy array."""
 
-    model_config = ConfigDict(validate_assignment=True, strict=False)
-
-    @staticmethod
-    def row(cls: type[TableSubclass], **kwargs) -> TableSubclass:  # type: ignore
-        arrayified_kwargs = {}
-        for field_name, field_value in cls.model_fields.items():
-            value = kwargs.pop(field_name)
-            if field_value.default_factory is None:
-                raise ValueError(
-                    "`Table` models should have default factories for their "
-                    "mutable empty columns."
-                )
-            default_array = field_value.default_factory()
-            if isinstance(default_array, np.ndarray):
-                arrayified_kwargs[field_name] = np.array(
-                    [value], dtype=default_array.dtype
-                )
-            elif issubclass(type(value), Enum) and isinstance(value, str):
-                arrayified_kwargs[field_name] = [value]
+    # You can use Table in 2 ways:
+    # 1. Table(**whatever_pva_gives_us) when pvi adds a Signal to a Device that is not
+    #    type hinted
+    # 2. MyTable(**whatever_pva_gives_us) where the Signal is type hinted
+    #
+    # For 1 we want extra="allow" so it is passed through as is. There are no base class
+    # fields, only "extra" fields, so they must be allowed. For 2 we want extra="forbid"
+    # so it is strictly checked against the BaseModel we are supplied.
+    model_config = ConfigDict(extra="allow")
+
+    @classmethod
+    def __init_subclass__(cls):
+        # But forbit extra in subclasses so it gets validated
+        cls.model_config = ConfigDict(validate_assignment=True, extra="forbid")
+        # Change fields to have the correct annotations
+        for k, anno in cls.__annotations__.items():
+            if get_origin(anno) is np.ndarray:
+                dtype = get_dtype(anno)
+                new_anno = Annotated[
+                    anno,
+                    NpArrayPydanticAnnotation.factory(
+                        data_type=dtype.type, dimensions=1, strict_data_typing=False
+                    ),
+                    Field(default_factory=_make_default_factory(dtype)),
+                ]
+            elif get_origin(anno) is Sequence:
+                new_anno = Annotated[anno, Field(default_factory=list)]
             else:
-                raise TypeError(
-                    "Row column should be numpy arrays or sequence of string `Enum`."
-                )
-        if kwargs:
-            raise TypeError(
-                f"Unexpected keyword arguments {kwargs.keys()} for {cls.__name__}."
-            )
-        return cls(**arrayified_kwargs)
+                raise TypeError(f"Cannot use annotation {anno} in a Table")
+            cls.__annotations__[k] = new_anno
 
     def __add__(self, right: TableSubclass) -> TableSubclass:
         """Concatenate the arrays in field values."""
@@ -64,83 +78,71 @@ class Table(BaseModel):
             }
         )
 
+    def __eq__(self, value: object) -> bool:
+        return super().__eq__(value)
+
     def numpy_dtype(self) -> np.dtype:
         dtype = []
-        for field_name, field_value in self.model_fields.items():
-            if np.ndarray in (
-                get_origin(field_value.annotation),
-                field_value.annotation,
-            ):
-                dtype.append((field_name, getattr(self, field_name).dtype))
+        for k, v in self:
+            if isinstance(v, np.ndarray):
+                dtype.append((k, v.dtype))
             else:
-                enum_type = get_args(field_value.annotation)[0]
-                assert issubclass(enum_type, Enum)
-                enum_values = [element.value for element in enum_type]
-                max_length_in_enum = max(len(value) for value in enum_values)
-                dtype.append((field_name, np.dtype(f"<U{max_length_in_enum}")))
-
+                # TODO: use np.dtypes.StringDType when we can use in structured arrays
+                # https://github.com/numpy/numpy/issues/25693
+                dtype.append((k, np.dtype("S40")))
         return np.dtype(dtype)
 
-    def numpy_table(self):
-        # It would be nice to be able to use np.transpose for this,
-        # but it defaults to the largest dtype for everything.
-        dtype = self.numpy_dtype()
-        transposed_list = [
-            np.array(tuple(row), dtype=dtype)
-            for row in zip(*self.numpy_columns(), strict=False)
-        ]
-        transposed = np.array(transposed_list, dtype=dtype)
-        return transposed
-
-    def numpy_columns(self) -> list[np.ndarray]:
-        """Columns in the table can be lists of string enums or numpy arrays.
-
-        This method returns the columns, converting the string enums to numpy arrays.
-        """
-
-        columns = []
-        for field_name, field_value in self.model_fields.items():
-            if np.ndarray in (
-                get_origin(field_value.annotation),
-                field_value.annotation,
+    def numpy_table(self, selection: slice | None = None) -> np.ndarray:
+        array = None
+        for k, v in self:
+            if selection:
+                v = v[selection]
+            if array is None:
+                array = np.empty(v.shape, dtype=self.numpy_dtype())
+            array[k] = v
+        assert array is not None
+        return array
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_array_dtypes(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            data_dict = data
+        elif isinstance(data, Table):
+            data_dict = data.model_dump()
+        else:
+            raise AssertionError(f"Cannot construct Table from {data}")
+        for field_name, field_value in cls.model_fields.items():
+            if (
+                get_origin(field_value.annotation) is np.ndarray
+                and field_value.annotation
+                and field_name in data_dict
             ):
-                columns.append(getattr(self, field_name))
-            else:
-                enum_type = get_args(field_value.annotation)[0]
-                assert issubclass(enum_type, Enum)
-                enum_values = [element.value for element in enum_type]
-                max_length_in_enum = max(len(value) for value in enum_values)
-                dtype = np.dtype(f"<U{max_length_in_enum}")
-
-                columns.append(
-                    np.array(
-                        [enum.value for enum in getattr(self, field_name)], dtype=dtype
-                    )
+                data_value = data_dict[field_name]
+                expected_dtype = get_dtype(field_value.annotation)
+                # Convert to correct dtype, but only if we don't lose precision
+                # as a result
+                cast_value = np.array(data_value).astype(expected_dtype)
+                assert np.array_equal(data_value, cast_value), (
+                    f"{field_name}: Cannot cast {data_value} to {expected_dtype} "
+                    "without losing precision"
                 )
-
-        return columns
+                data_dict[field_name] = cast_value
+        return data_dict
 
     @model_validator(mode="after")
-    def validate_arrays(self) -> "Table":
-        first_length = len(next(iter(self))[1])
-        assert all(
-            len(field_value) == first_length for _, field_value in self
-        ), "Rows should all be of equal size."
-
-        if not all(
-            # Checks if the values are numpy subtypes if the array is a numpy array,
-            # or if the value is a string enum.
-            np.issubdtype(getattr(self, field_name).dtype, default_array.dtype)
-            if isinstance(
-                default_array := self.model_fields[field_name].default_factory(),  # type: ignore
-                np.ndarray,
-            )
-            else issubclass(get_args(field_value.annotation)[0], Enum)
-            for field_name, field_value in self.model_fields.items()
-        ):
-            raise ValueError(
-                f"Cannot construct a `{type(self).__name__}`, "
-                "some rows have incorrect types."
-            )
-
+    def validate_lengths(self) -> Table:
+        lengths: dict[int, set[str]] = {}
+        for field_name, field_value in self:
+            lengths.setdefault(len(field_value), set()).add(field_name)
+        assert len(lengths) <= 1, f"Columns should be same length, got {lengths=}"
         return self
+
+    def __len__(self) -> int:
+        return len(next(iter(self))[1])
+
+    def __getitem__(self, item: int | slice) -> np.ndarray:
+        if isinstance(item, int):
+            return self.numpy_table(slice(item, item + 1))
+        else:
+            return self.numpy_table(item)

ophyd_async/core/_utils.py

@@ -2,23 +2,34 @@ from __future__ import annotations
 
 import asyncio
 import logging
-from collections.abc import Awaitable, Callable, Iterable
+from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
 from dataclasses import dataclass
-from typing import Generic, Literal, ParamSpec, TypeVar, get_origin
+from enum import Enum, EnumMeta
+from typing import Any, Generic, Literal, ParamSpec, TypeVar, get_args, get_origin
+from unittest.mock import Mock
 
 import numpy as np
-from bluesky.protocols import Reading
-from pydantic import BaseModel
 
 T = TypeVar("T")
 P = ParamSpec("P")
 Callback = Callable[[T], None]
-
-#: A function that will be called with the Reading and value when the
-#: monitor updates
-ReadingValueCallback = Callable[[Reading, T], None]
 DEFAULT_TIMEOUT = 10.0
-ErrorText = str | dict[str, Exception]
+ErrorText = str | Mapping[str, Exception]
+
+
+class StrictEnum(str, Enum):
+    """All members should exist in the Backend, and there will be no extras"""
+
+
+class SubsetEnumMeta(EnumMeta):
+    def __call__(self, value, *args, **kwargs):  # type: ignore
+        if isinstance(value, str) and not isinstance(value, self):
+            return value
+        return super().__call__(value, *args, **kwargs)
+
+
+class SubsetEnum(StrictEnum, metaclass=SubsetEnumMeta):
+    """All members should exist in the Backend, but there may be extras"""
 
 
 CALCULATE_TIMEOUT = "CALCULATE_TIMEOUT"
@@ -51,6 +62,13 @@ class NotConnected(Exception):
 
         self._errors = errors
 
+    @property
+    def sub_errors(self) -> Mapping[str, Exception]:
+        if isinstance(self._errors, dict):
+            return self._errors.copy()
+        else:
+            return {}
+
     def _format_sub_errors(self, name: str, error: Exception, indent="") -> str:
         if isinstance(error, NotConnected):
             error_txt = ":" + error.format_error_string(indent + self._indent_width)
@@ -81,6 +99,19 @@ class NotConnected(Exception):
     def __str__(self) -> str:
         return self.format_error_string(indent="")
 
+    @classmethod
+    def with_other_exceptions_logged(
+        cls, exceptions: Mapping[str, Exception]
+    ) -> NotConnected:
+        for name, exception in exceptions.items():
+            if not isinstance(exception, NotConnected):
+                logging.exception(
+                    f"device `{name}` raised unexpected exception "
+                    f"{type(exception).__name__}",
+                    exc_info=exception,
+                )
+        return NotConnected(exceptions)
+
 
 @dataclass(frozen=True)
 class WatcherUpdate(Generic[T]):
@@ -102,24 +133,41 @@ async def wait_for_connection(**coros: Awaitable[None]):
 
     Expected kwargs should be a mapping of names to coroutine tasks to execute.
     """
-    results = await asyncio.gather(*coros.values(), return_exceptions=True)
-    exceptions = {}
-
-    for name, result in zip(coros, results, strict=False):
-        if isinstance(result, Exception):
-            exceptions[name] = result
-            if not isinstance(result, NotConnected):
-                logging.exception(
-                    f"device `{name}` raised unexpected exception "
-                    f"{type(result).__name__}",
-                    exc_info=result,
-                )
+    exceptions: dict[str, Exception] = {}
+    if len(coros) == 1:
+        # Single device optimization
+        name, coro = coros.popitem()
+        try:
+            await coro
+        except Exception as e:
+            exceptions[name] = e
+    else:
+        # Use gather to connect in parallel
+        results = await asyncio.gather(*coros.values(), return_exceptions=True)
+        for name, result in zip(coros, results, strict=False):
+            if isinstance(result, Exception):
+                exceptions[name] = result
 
     if exceptions:
-        raise NotConnected(exceptions)
+        raise NotConnected.with_other_exceptions_logged(exceptions)
 
 
-def get_dtype(typ: type) -> np.dtype | None:
+def get_dtype(datatype: type) -> np.dtype:
+    """Get the runtime dtype from a numpy ndarray type annotation
+
+    >>> from ophyd_async.core import Array1D
+    >>> import numpy as np
+    >>> get_dtype(Array1D[np.int8])
+    dtype('int8')
+    """
+    if not get_origin(datatype) == np.ndarray:
+        raise TypeError(f"Expected Array1D[dtype], got {datatype}")
+    # datatype = numpy.ndarray[typing.Any, numpy.dtype[numpy.float64]]
+    # so extract numpy.float64 from it
+    return np.dtype(get_args(get_args(datatype)[1])[0])
+
+
+def get_enum_cls(datatype: type | None) -> type[StrictEnum] | None:
     """Get the runtime dtype from a numpy ndarray type annotation
 
     >>> import numpy.typing as npt
@@ -127,11 +175,15 @@ def get_dtype(typ: type) -> np.dtype | None:
     >>> get_dtype(npt.NDArray[np.int8])
     dtype('int8')
     """
-    if getattr(typ, "__origin__", None) == np.ndarray:
-        # datatype = numpy.ndarray[typing.Any, numpy.dtype[numpy.float64]]
-        # so extract numpy.float64 from it
-        return np.dtype(typ.__args__[1].__args__[0])  # type: ignore
-    return None
+    if get_origin(datatype) is Sequence:
+        datatype = get_args(datatype)[0]
+    if datatype and issubclass(datatype, Enum):
+        if not issubclass(datatype, StrictEnum):
+            raise TypeError(
+                f"{datatype} should inherit from .SubsetEnum "
+                "or ophyd_async.core.StrictEnum"
+            )
+        return datatype
 
 
 def get_unique(values: dict[str, T], types: str) -> T:
@@ -187,7 +239,66 @@ def in_micros(t: float) -> int:
     return int(np.ceil(t * 1e6))
 
 
-def is_pydantic_model(datatype) -> bool:
-    while origin := get_origin(datatype):
-        datatype = origin
-    return datatype and issubclass(datatype, BaseModel)
+def get_origin_class(annotatation: Any) -> type | None:
+    origin = get_origin(annotatation) or annotatation
+    if isinstance(origin, type):
+        return origin
+
+
+class Reference(Generic[T]):
+    """Hide an object behind a reference.
+
+    Used to opt out of the naming/parent-child relationship of `Device`.
+
+    For example::
+
+        class DeviceWithRefToSignal(Device):
+            def __init__(self, signal: SignalRW[int]):
+                self.signal_ref = Reference(signal)
+                super().__init__()
+
+            def set(self, value) -> AsyncStatus:
+                return self.signal_ref().set(value + 1)
+
+    """
+
+    def __init__(self, obj: T):
+        self._obj = obj
+
+    def __call__(self) -> T:
+        return self._obj
+
+
+class LazyMock:
+    """A lazily created Mock to be used when connecting in mock mode.
+
+    Creating Mocks is reasonably expensive when each Device (and Signal)
+    requires its own, and the tree is only used when ``Signal.set()`` is
+    called. This class allows a tree of lazily connected Mocks to be
+    constructed so that when the leaf is created, so are its parents.
+    Any calls to the child are then accessible from the parent mock.
+
+    >>> parent = LazyMock()
+    >>> child = parent.child("child")
+    >>> child_mock = child()
+    >>> child_mock()  # doctest: +ELLIPSIS
+    <Mock name='mock.child()' id='...'>
+    >>> parent_mock = parent()
+    >>> parent_mock.mock_calls
+    [call.child()]
+    """
+
+    def __init__(self, name: str = "", parent: LazyMock | None = None) -> None:
+        self.parent = parent
+        self.name = name
+        self._mock: Mock | None = None
+
+    def child(self, name: str) -> LazyMock:
+        return LazyMock(name, self)
+
+    def __call__(self) -> Mock:
+        if self._mock is None:
+            self._mock = Mock(spec=object)
+            if self.parent is not None:
+                self.parent().attach_mock(self._mock, self.name)
+        return self._mock

ophyd_async/epics/adaravis/_aravis_controller.py

@@ -2,12 +2,12 @@ import asyncio
 from typing import Literal
 
 from ophyd_async.core import (
+    AsyncStatus,
     DetectorController,
     DetectorTrigger,
     TriggerInfo,
     set_and_wait_for_value,
 )
-from ophyd_async.core._status import AsyncStatus
 from ophyd_async.epics import adcore
 
 from ._aravis_io import AravisDriverIO, AravisTriggerMode, AravisTriggerSource
@@ -69,7 +69,7 @@ class AravisController(DetectorController):
                 f"use {trigger}"
             )
         if trigger == DetectorTrigger.internal:
-            return AravisTriggerMode.off, "Freerun"
+            return AravisTriggerMode.off, AravisTriggerSource.freerun
         else:
             return (AravisTriggerMode.on, f"Line{self.gpio_number}")  # type: ignore
 

ophyd_async/epics/adaravis/_aravis_io.py

@@ -1,11 +1,9 @@
-from enum import Enum
-
-from ophyd_async.core import SubsetEnum
+from ophyd_async.core import StrictEnum, SubsetEnum
 from ophyd_async.epics import adcore
-from ophyd_async.epics.signal import epics_signal_rw_rbv
+from ophyd_async.epics.core import epics_signal_rw_rbv
 
 
-class AravisTriggerMode(str, Enum):
+class AravisTriggerMode(StrictEnum):
     """GigEVision GenICAM standard: on=externally triggered"""
 
     on = "On"
@@ -19,7 +17,11 @@ class AravisTriggerMode(str, Enum):
     To prevent requiring one Enum class per possible configuration, we set as this Enum
     but read from the underlying signal as a str.
     """
-AravisTriggerSource = SubsetEnum["Freerun", "Line1"]
+
+
+class AravisTriggerSource(SubsetEnum):
+    freerun = "Freerun"
+    line1 = "Line1"
 
 
 class AravisDriverIO(adcore.ADBaseIO):

ophyd_async/epics/adcore/_core_io.py

@@ -1,7 +1,5 @@
-from enum import Enum
-
-from ophyd_async.core import Device
-from ophyd_async.epics.signal import (
+from ophyd_async.core import Device, StrictEnum
+from ophyd_async.epics.core import (
     epics_signal_r,
     epics_signal_rw,
     epics_signal_rw_rbv,
@@ -10,7 +8,7 @@ from ophyd_async.epics.signal import (
 from ._utils import ADBaseDataType, FileWriteMode, ImageMode
 
 
-class Callback(str, Enum):
+class Callback(StrictEnum):
     Enable = "Enable"
     Disable = "Disable"
 
@@ -68,7 +66,7 @@ class NDPluginStatsIO(NDPluginBaseIO):
         super().__init__(prefix, name)
 
 
-class DetectorState(str, Enum):
+class DetectorState(StrictEnum):
     """
     Default set of states of an AreaDetector driver.
     See definition in ADApp/ADSrc/ADDriver.h in https://github.com/areaDetector/ADCore
@@ -100,7 +98,7 @@ class ADBaseIO(NDArrayBaseIO):
         super().__init__(prefix, name=name)
 
 
-class Compression(str, Enum):
+class Compression(StrictEnum):
     none = "None"
     nbit = "N-bit"
     szip = "szip"

ophyd_async/epics/adcore/_core_logic.py

@@ -91,7 +91,9 @@ async def start_acquiring_driver_and_ensure_status(
         subsequent raising (if applicable) due to detector state.
     """
 
-    status = await set_and_wait_for_value(driver.acquire, True, timeout=timeout)
+    status = await set_and_wait_for_value(
+        driver.acquire, True, timeout=timeout, wait_for_set_completion=False
+    )
 
     async def complete_acquisition() -> None:
         """NOTE: possible race condition here between the callback from

ophyd_async/epics/adcore/_hdf_writer.py

@@ -134,9 +134,9 @@ class ADHDFWriter(DetectorWriter):
         describe = {
             ds.data_key: DataKey(
                 source=self.hdf.full_file_name.source,
-                shape=outer_shape + tuple(ds.shape),
+                shape=list(outer_shape + tuple(ds.shape)),
                 dtype="array" if ds.shape else "number",
-                dtype_numpy=ds.dtype_numpy,  # type: ignore
+                dtype_numpy=ds.dtype_numpy,
                 external="STREAM:",
             )
             for ds in self._datasets