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

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 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")
 
@@ -17,34 +22,38 @@ def _concat(value1, value2):
 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=lambda dtype=dtype: np.array([], dtype=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 +73,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,25 +2,43 @@ from __future__ import annotations
 
 import asyncio
 import logging
-from collections.abc import Awaitable, Callable, Iterable
+from collections.abc import Awaitable, Callable, Iterable, 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,
+)
 
 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]
 
 
+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"
 """Sentinel used to implement ``myfunc(timeout=CalculateTimeout)``
 
@@ -119,7 +137,22 @@ async def wait_for_connection(**coros: Awaitable[None]):
         raise NotConnected(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 +160,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 +224,31 @@ 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

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/_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

ophyd_async/epics/adcore/_single_trigger.py

@@ -3,13 +3,8 @@ from collections.abc import Sequence
 
 from bluesky.protocols import Triggerable
 
-from ophyd_async.core import (
-    AsyncStatus,
-    ConfigSignal,
-    HintedSignal,
-    SignalR,
-    StandardReadable,
-)
+from ophyd_async.core import AsyncStatus, SignalR, StandardReadable
+from ophyd_async.core import StandardReadableFormat as Format
 
 from ._core_io import ADBaseIO, NDPluginBaseIO
 from ._utils import ImageMode
@@ -28,10 +23,10 @@ class SingleTriggerDetector(StandardReadable, Triggerable):
 
         self.add_readables(
             [self.drv.array_counter, *read_uncached],
-            wrapper=HintedSignal.uncached,
+            Format.HINTED_UNCACHED_SIGNAL,
         )
 
-        self.add_readables([self.drv.acquire_time], wrapper=ConfigSignal)
+        self.add_readables([self.drv.acquire_time], Format.CONFIG_SIGNAL)
 
         super().__init__(name=name)
 

ophyd_async/epics/adcore/_utils.py

@@ -1,11 +1,16 @@
 from dataclasses import dataclass
-from enum import Enum
 
-from ophyd_async.core import DEFAULT_TIMEOUT, SignalRW, T, wait_for_value
-from ophyd_async.core._signal import SignalR
+from ophyd_async.core import (
+    DEFAULT_TIMEOUT,
+    SignalDatatypeT,
+    SignalR,
+    SignalRW,
+    StrictEnum,
+    wait_for_value,
+)
 
 
-class ADBaseDataType(str, Enum):
+class ADBaseDataType(StrictEnum):
     Int8 = "Int8"
     UInt8 = "UInt8"
     Int16 = "Int16"
@@ -73,25 +78,25 @@ def convert_param_dtype_to_np(datatype: str) -> str:
     return np_datatype
 
 
-class FileWriteMode(str, Enum):
+class FileWriteMode(StrictEnum):
     single = "Single"
     capture = "Capture"
     stream = "Stream"
 
 
-class ImageMode(str, Enum):
+class ImageMode(StrictEnum):
     single = "Single"
     multiple = "Multiple"
     continuous = "Continuous"
 
 
-class NDAttributeDataType(str, Enum):
+class NDAttributeDataType(StrictEnum):
     INT = "INT"
     DOUBLE = "DOUBLE"
     STRING = "STRING"
 
 
-class NDAttributePvDbrType(str, Enum):
+class NDAttributePvDbrType(StrictEnum):
     DBR_SHORT = "DBR_SHORT"
     DBR_ENUM = "DBR_ENUM"
     DBR_INT = "DBR_INT"
@@ -122,8 +127,8 @@ class NDAttributeParam:
 
 
 async def stop_busy_record(
-    signal: SignalRW[T],
-    value: T,
+    signal: SignalRW[SignalDatatypeT],
+    value: SignalDatatypeT,
     timeout: float = DEFAULT_TIMEOUT,
     status_timeout: float | None = None,
 ) -> None:

ophyd_async/epics/adkinetix/__init__.py

@@ -1,9 +1,10 @@
 from ._kinetix import KinetixDetector
 from ._kinetix_controller import KinetixController
-from ._kinetix_io import KinetixDriverIO
+from ._kinetix_io import KinetixDriverIO, KinetixTriggerMode
 
 __all__ = [
     "KinetixDetector",
     "KinetixController",
     "KinetixDriverIO",
+    "KinetixTriggerMode",
 ]

ophyd_async/epics/adkinetix/_kinetix_controller.py

@@ -1,8 +1,11 @@
 import asyncio
 
-from ophyd_async.core import DetectorController, DetectorTrigger
-from ophyd_async.core._detector import TriggerInfo
-from ophyd_async.core._status import AsyncStatus
+from ophyd_async.core import (
+    AsyncStatus,
+    DetectorController,
+    DetectorTrigger,
+    TriggerInfo,
+)
 from ophyd_async.epics import adcore
 
 from ._kinetix_io import KinetixDriverIO, KinetixTriggerMode

ophyd_async/epics/adkinetix/_kinetix_io.py

@@ -1,16 +1,15 @@
-from enum import Enum
-
+from ophyd_async.core import StrictEnum
 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 KinetixTriggerMode(str, Enum):
+class KinetixTriggerMode(StrictEnum):
     internal = "Internal"
     edge = "Rising Edge"
     gate = "Exp. Gate"
 
 
-class KinetixReadoutMode(str, Enum):
+class KinetixReadoutMode(StrictEnum):
     sensitivity = 1
     speed = 2
     dynamic_range = 3

ophyd_async/epics/adpilatus/_pilatus_controller.py

@@ -2,12 +2,12 @@ import asyncio
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
+    AsyncStatus,
     DetectorController,
     DetectorTrigger,
+    TriggerInfo,
     wait_for_value,
 )
-from ophyd_async.core._detector import TriggerInfo
-from ophyd_async.core._status import AsyncStatus
 from ophyd_async.epics import adcore
 
 from ._pilatus_io import PilatusDriverIO, PilatusTriggerMode

ophyd_async/epics/adpilatus/_pilatus_io.py

@@ -1,10 +1,9 @@
-from enum import Enum
-
+from ophyd_async.core import StrictEnum
 from ophyd_async.epics import adcore
-from ophyd_async.epics.signal import epics_signal_r, epics_signal_rw_rbv
+from ophyd_async.epics.core import epics_signal_r, epics_signal_rw_rbv
 
 
-class PilatusTriggerMode(str, Enum):
+class PilatusTriggerMode(StrictEnum):
     internal = "Internal"
     ext_enable = "Ext. Enable"
     ext_trigger = "Ext. Trigger"

ophyd_async/epics/adsimdetector/_sim_controller.py

@@ -2,11 +2,11 @@ import asyncio
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
+    AsyncStatus,
     DetectorController,
     DetectorTrigger,
+    TriggerInfo,
 )
-from ophyd_async.core._detector import TriggerInfo
-from ophyd_async.core._status import AsyncStatus
 from ophyd_async.epics import adcore
 
 

ophyd_async/epics/advimba/__init__.py

@@ -1,9 +1,12 @@
 from ._vimba import VimbaDetector
 from ._vimba_controller import VimbaController
-from ._vimba_io import VimbaDriverIO
+from ._vimba_io import VimbaDriverIO, VimbaExposeOutMode, VimbaOnOff, VimbaTriggerSource
 
 __all__ = [
     "VimbaDetector",
     "VimbaController",
     "VimbaDriverIO",
+    "VimbaExposeOutMode",
+    "VimbaOnOff",
+    "VimbaTriggerSource",
 ]