ophyd-async 0.5.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

ophyd_async/epics/signal/_aioca.py

@@ -1,9 +1,11 @@
+import inspect
 import logging
 import sys
+from collections.abc import Sequence
 from dataclasses import dataclass
 from enum import Enum
 from math import isnan, nan
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import Any, get_origin
 
 import numpy as np
 from aioca import (
@@ -17,13 +19,16 @@ from aioca import (
     caput,
 )
 from aioca.types import AugmentedValue, Dbr, Format
-from bluesky.protocols import DataKey, Dtype, Reading
+from bluesky.protocols import Reading
 from epicscorelibs.ca import dbr
+from event_model import DataKey
+from event_model.documents.event_descriptor import Dtype
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
     NotConnected,
     ReadingValueCallback,
+    RuntimeSubsetEnum,
     SignalBackend,
     T,
     get_dtype,
@@ -33,7 +38,7 @@ from ophyd_async.core import (
 
 from ._common import LimitPair, Limits, common_meta, get_supported_values
 
-dbr_to_dtype: Dict[Dbr, Dtype] = {
+dbr_to_dtype: dict[Dbr, Dtype] = {
     dbr.DBR_STRING: "string",
     dbr.DBR_SHORT: "integer",
     dbr.DBR_FLOAT: "number",
@@ -46,8 +51,8 @@ dbr_to_dtype: Dict[Dbr, Dtype] = {
 def _data_key_from_augmented_value(
     value: AugmentedValue,
     *,
-    choices: Optional[List[str]] = None,
-    dtype: Optional[Dtype] = None,
+    choices: list[str] | None = None,
+    dtype: Dtype | None = None,
 ) -> DataKey:
     """Use the return value of get with FORMAT_CTRL to construct a DataKey
     describing the signal. See docstring of AugmentedValue for expected
@@ -65,14 +70,15 @@ def _data_key_from_augmented_value(
     assert value.ok, f"Error reading {source}: {value}"
 
     scalar = value.element_count == 1
-    dtype = dtype or dbr_to_dtype[value.datatype]
+    dtype = dtype or dbr_to_dtype[value.datatype]  # type: ignore
 
     dtype_numpy = np.dtype(dbr.DbrCodeToType[value.datatype].dtype).descr[0][1]
 
     d = DataKey(
         source=source,
         dtype=dtype if scalar else "array",
-        dtype_numpy=dtype_numpy,
+        # Ignore until https://github.com/bluesky/event-model/issues/308
+        dtype_numpy=dtype_numpy,  # type: ignore
         # strictly value.element_count >= len(value)
         shape=[] if scalar else [len(value)],
     )
@@ -82,10 +88,10 @@ def _data_key_from_augmented_value(
             d[key] = attr
 
     if choices is not None:
-        d["choices"] = choices
+        d["choices"] = choices  # type: ignore
 
     if limits := _limits_from_augmented_value(value):
-        d["limits"] = limits
+        d["limits"] = limits  # type: ignore
 
     return d
 
@@ -108,8 +114,8 @@ def _limits_from_augmented_value(value: AugmentedValue) -> Limits:
 
 @dataclass
 class CaConverter:
-    read_dbr: Optional[Dbr]
-    write_dbr: Optional[Dbr]
+    read_dbr: Dbr | None
+    write_dbr: Dbr | None
 
     def write_value(self, value) -> Any:
         return value
@@ -118,9 +124,9 @@ class CaConverter:
         # for channel access ca_xxx classes, this
         # invokes __pos__ operator to return an instance of
         # the builtin base class
-        return +value
+        return +value  # type: ignore
 
-    def reading(self, value: AugmentedValue):
+    def reading(self, value: AugmentedValue) -> Reading:
         return {
             "value": self.value(value),
             "timestamp": value.timestamp,
@@ -155,14 +161,14 @@ class CaEnumConverter(CaConverter):
 
     choices: dict[str, str]
 
-    def write_value(self, value: Union[Enum, str]):
+    def write_value(self, value: Enum | str):
         if isinstance(value, Enum):
             return value.value
         else:
             return value
 
     def value(self, value: AugmentedValue):
-        return self.choices[value]
+        return self.choices[value]  # type: ignore
 
     def get_datakey(self, value: AugmentedValue) -> DataKey:
         # Sometimes DBR_TYPE returns as String, must pass choices still
@@ -184,7 +190,7 @@ class DisconnectedCaConverter(CaConverter):
 
 
 def make_converter(
-    datatype: Optional[Type], values: Dict[str, AugmentedValue]
+    datatype: type | None, values: dict[str, AugmentedValue]
 ) -> CaConverter:
     pv = list(values)[0]
     pv_dbr = get_unique({k: v.datatype for k, v in values.items()}, "datatypes")
@@ -200,7 +206,7 @@ def make_converter(
                 raise TypeError(f"{pv} has type [str] not {datatype.__name__}")
         return CaArrayConverter(pv_dbr, None)
     elif is_array:
-        pv_dtype = get_unique({k: v.dtype for k, v in values.items()}, "dtypes")
+        pv_dtype = get_unique({k: v.dtype for k, v in values.items()}, "dtypes")  # type: ignore
         # This is an array
         if datatype:
             # Check we wanted an array of this type
@@ -209,9 +215,10 @@ def make_converter(
                 raise TypeError(f"{pv} has type [{pv_dtype}] not {datatype.__name__}")
             if dtype != pv_dtype:
                 raise TypeError(f"{pv} has type [{pv_dtype}] not [{dtype}]")
-        return CaArrayConverter(pv_dbr, None)
+        return CaArrayConverter(pv_dbr, None)  # type: ignore
     elif pv_dbr == dbr.DBR_ENUM and datatype is bool:
-        # Database can't do bools, so are often representated as enums, CA can do int
+        # Database can't do bools, so are often representated as enums,
+        # CA can do int
         pv_choices_len = get_unique(
             {k: len(v.enums) for k, v in values.items()}, "number of choices"
         )
@@ -240,7 +247,7 @@ def make_converter(
                     f"{pv} has type {type(value).__name__.replace('ca_', '')} "
                     + f"not {datatype.__name__}"
                 )
-        return CaConverter(pv_dbr, None)
+    return CaConverter(pv_dbr, None)  # type: ignore
 
 
 _tried_pyepics = False
@@ -256,13 +263,36 @@ def _use_pyepics_context_if_imported():
 
 
 class CaSignalBackend(SignalBackend[T]):
-    def __init__(self, datatype: Optional[Type[T]], read_pv: str, write_pv: str):
+    _ALLOWED_DATATYPES = (
+        bool,
+        int,
+        float,
+        str,
+        Sequence,
+        Enum,
+        RuntimeSubsetEnum,
+        np.ndarray,
+    )
+
+    @classmethod
+    def datatype_allowed(cls, dtype: Any) -> bool:
+        stripped_origin = get_origin(dtype) or dtype
+        if dtype is None:
+            return True
+
+        return inspect.isclass(stripped_origin) and issubclass(
+            stripped_origin, cls._ALLOWED_DATATYPES
+        )
+
+    def __init__(self, datatype: type[T] | None, read_pv: str, write_pv: str):
         self.datatype = datatype
+        if not CaSignalBackend.datatype_allowed(self.datatype):
+            raise TypeError(f"Given datatype {self.datatype} unsupported in CA.")
         self.read_pv = read_pv
         self.write_pv = write_pv
-        self.initial_values: Dict[str, AugmentedValue] = {}
+        self.initial_values: dict[str, AugmentedValue] = {}
         self.converter: CaConverter = DisconnectedCaConverter(None, None)
-        self.subscription: Optional[Subscription] = None
+        self.subscription: Subscription | None = None
 
     def source(self, name: str):
         return f"ca://{self.read_pv}"
@@ -289,7 +319,7 @@ class CaSignalBackend(SignalBackend[T]):
             await self._store_initial_value(self.read_pv, timeout=timeout)
         self.converter = make_converter(self.datatype, self.initial_values)
 
-    async def put(self, value: Optional[T], wait=True, timeout=None):
+    async def put(self, value: T | None, wait=True, timeout=None):
         if value is None:
             write_value = self.initial_values[self.write_pv]
         else:
@@ -331,7 +361,7 @@ class CaSignalBackend(SignalBackend[T]):
         )
         return self.converter.value(value)
 
-    def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:
+    def set_callback(self, callback: ReadingValueCallback[T] | None) -> None:
         if callback:
             assert (
                 not self.subscription

ophyd_async/epics/signal/_common.py

@@ -1,6 +1,5 @@
 import inspect
 from enum import Enum
-from typing import Dict, Optional, Tuple, Type
 
 from typing_extensions import TypedDict
 
@@ -16,9 +15,6 @@ class LimitPair(TypedDict):
     high: float | None
     low: float | None
 
-    def __bool__(self) -> bool:
-        return self.low is None and self.high is None
-
 
 class Limits(TypedDict):
     alarm: LimitPair
@@ -26,15 +22,12 @@ class Limits(TypedDict):
     display: LimitPair
     warning: LimitPair
 
-    def __bool__(self) -> bool:
-        return any(self.alarm, self.control, self.display, self.warning)
-
 
 def get_supported_values(
     pv: str,
-    datatype: Optional[Type[str]],
-    pv_choices: Tuple[str, ...],
-) -> Dict[str, str]:
+    datatype: type[str] | None,
+    pv_choices: tuple[str, ...],
+) -> dict[str, str]:
     if inspect.isclass(datatype) and issubclass(datatype, RuntimeSubsetEnum):
         if not set(datatype.choices).issubset(set(pv_choices)):
             raise TypeError(

ophyd_async/epics/signal/_epics_transport.py

@@ -4,22 +4,25 @@ from __future__ import annotations
 
 from enum import Enum
 
+
+def _make_unavailable_class(error: Exception) -> type:
+    class TransportNotAvailable:
+        def __init__(*args, **kwargs):
+            raise NotImplementedError("Transport not available") from error
+
+    return TransportNotAvailable
+
+
 try:
     from ._aioca import CaSignalBackend
 except ImportError as ca_error:
-
-    class CaSignalBackend:  # type: ignore
-        def __init__(*args, ca_error=ca_error, **kwargs):
-            raise NotImplementedError("CA support not available") from ca_error
+    CaSignalBackend = _make_unavailable_class(ca_error)
 
 
 try:
     from ._p4p import PvaSignalBackend
 except ImportError as pva_error:
-
-    class PvaSignalBackend:  # type: ignore
-        def __init__(*args, pva_error=pva_error, **kwargs):
-            raise NotImplementedError("PVA support not available") from pva_error
+    PvaSignalBackend = _make_unavailable_class(pva_error)
 
 
 class _EpicsTransport(Enum):

ophyd_async/epics/signal/_p4p.py

@@ -3,14 +3,19 @@ import atexit
 import inspect
 import logging
 import time
+from collections.abc import Sequence
 from dataclasses import dataclass
 from enum import Enum
 from math import isnan, nan
-from typing import Any, Dict, List, Optional, Sequence, Type, Union
+from typing import Any, get_origin
 
-from bluesky.protocols import DataKey, Dtype, Reading
+import numpy as np
+from bluesky.protocols import Reading
+from event_model import DataKey
+from event_model.documents.event_descriptor import Dtype
 from p4p import Value
 from p4p.client.asyncio import Context, Subscription
+from pydantic import BaseModel
 
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
@@ -21,13 +26,14 @@ from ophyd_async.core import (
     T,
     get_dtype,
     get_unique,
+    is_pydantic_model,
     wait_for_connection,
 )
 
 from ._common import LimitPair, Limits, common_meta, get_supported_values
 
 # https://mdavidsaver.github.io/p4p/values.html
-specifier_to_dtype: Dict[str, Dtype] = {
+specifier_to_dtype: dict[str, Dtype] = {
     "?": "integer",  # bool
     "b": "integer",  # int8
     "B": "integer",  # uint8
@@ -42,7 +48,7 @@ specifier_to_dtype: Dict[str, Dtype] = {
     "s": "string",
 }
 
-specifier_to_np_dtype: Dict[str, str] = {
+specifier_to_np_dtype: dict[str, str] = {
     "?": "<i2",  # bool
     "b": "|i1",  # int8
     "B": "|u1",  # uint8
@@ -62,9 +68,9 @@ def _data_key_from_value(
     source: str,
     value: Value,
     *,
-    shape: Optional[list[int]] = None,
-    choices: Optional[list[str]] = None,
-    dtype: Optional[Dtype] = None,
+    shape: list[int] | None = None,
+    choices: list[str] | None = None,
+    dtype: Dtype | None = None,
 ) -> DataKey:
     """
     Args:
@@ -105,7 +111,8 @@ def _data_key_from_value(
     d = DataKey(
         source=source,
         dtype=dtype,
-        dtype_numpy=dtype_numpy,
+        # type ignore until https://github.com/bluesky/event-model/issues/308
+        dtype_numpy=dtype_numpy,  # type: ignore
         shape=shape,
     )
     if display_data is not None:
@@ -115,10 +122,12 @@ def _data_key_from_value(
                 d[key] = attr
 
     if choices is not None:
-        d["choices"] = choices
+        # type ignore until https://github.com/bluesky/event-model/issues/309
+        d["choices"] = choices  # type: ignore
 
     if limits := _limits_from_value(value):
-        d["limits"] = limits
+        # type ignore until https://github.com/bluesky/event-model/issues/309
+        d["limits"] = limits  # type: ignore
 
     return d
 
@@ -149,7 +158,7 @@ class PvaConverter:
     def value(self, value):
         return value["value"]
 
-    def reading(self, value):
+    def reading(self, value) -> Reading:
         ts = value["timeStamp"]
         sv = value["alarm"]["severity"]
         return {
@@ -161,13 +170,13 @@ class PvaConverter:
     def get_datakey(self, source: str, value) -> DataKey:
         return _data_key_from_value(source, value)
 
-    def metadata_fields(self) -> List[str]:
+    def metadata_fields(self) -> list[str]:
         """
         Fields to request from PVA for metadata.
         """
         return ["alarm", "timeStamp"]
 
-    def value_fields(self) -> List[str]:
+    def value_fields(self) -> list[str]:
         """
         Fields to request from PVA for the value.
         """
@@ -182,11 +191,11 @@ class PvaArrayConverter(PvaConverter):
 
 
 class PvaNDArrayConverter(PvaConverter):
-    def metadata_fields(self) -> List[str]:
+    def metadata_fields(self) -> list[str]:
         return super().metadata_fields() + ["dimension"]
 
-    def _get_dimensions(self, value) -> List[int]:
-        dimensions: List[Value] = value["dimension"]
+    def _get_dimensions(self, value) -> list[int]:
+        dimensions: list[Value] = value["dimension"]
         dims = [dim.size for dim in dimensions]
         # Note: dimensions in NTNDArray are in fortran-like order
         # with first index changing fastest.
@@ -221,7 +230,7 @@ class PvaEnumConverter(PvaConverter):
     def __init__(self, choices: dict[str, str]):
         self.choices = tuple(choices.values())
 
-    def write_value(self, value: Union[Enum, str]):
+    def write_value(self, value: Enum | str):
         if isinstance(value, Enum):
             return value.value
         else:
@@ -250,11 +259,24 @@ class PvaTableConverter(PvaConverter):
 
     def get_datakey(self, source: str, value) -> DataKey:
         # This is wrong, but defer until we know how to actually describe a table
-        return _data_key_from_value(source, value, dtype="object")
+        return _data_key_from_value(source, value, dtype="object")  # type: ignore
+
+
+class PvaPydanticModelConverter(PvaConverter):
+    def __init__(self, datatype: BaseModel):
+        self.datatype = datatype
+
+    def value(self, value: Value):
+        return self.datatype(**value.todict())  # type: ignore
+
+    def write_value(self, value: BaseModel | dict[str, Any]):
+        if isinstance(value, self.datatype):  # type: ignore
+            return value.model_dump(mode="python")  # type: ignore
+        return value
 
 
 class PvaDictConverter(PvaConverter):
-    def reading(self, value):
+    def reading(self, value) -> Reading:
         ts = time.time()
         value = value.todict()
         # Alarm severity is vacuously 0 for a table
@@ -266,13 +288,13 @@ class PvaDictConverter(PvaConverter):
     def get_datakey(self, source: str, value) -> DataKey:
         raise NotImplementedError("Describing Dict signals not currently supported")
 
-    def metadata_fields(self) -> List[str]:
+    def metadata_fields(self) -> list[str]:
         """
         Fields to request from PVA for metadata.
         """
         return []
 
-    def value_fields(self) -> List[str]:
+    def value_fields(self) -> list[str]:
         """
         Fields to request from PVA for the value.
         """
@@ -284,7 +306,7 @@ class DisconnectedPvaConverter(PvaConverter):
         raise NotImplementedError("No PV has been set as connect() has not been called")
 
 
-def make_converter(datatype: Optional[Type], values: Dict[str, Any]) -> PvaConverter:
+def make_converter(datatype: type | None, values: dict[str, Any]) -> PvaConverter:
     pv = list(values)[0]
     typeid = get_unique({k: v.getID() for k, v in values.items()}, "typeids")
     typ = get_unique(
@@ -333,7 +355,7 @@ def make_converter(datatype: Optional[Type], values: Dict[str, Any]) -> PvaConve
             and issubclass(datatype, RuntimeSubsetEnum)
         ):
             return PvaEnumConverter(
-                get_supported_values(pv, datatype, datatype.choices)
+                get_supported_values(pv, datatype, datatype.choices)  # type: ignore
             )
         elif datatype and not issubclass(typ, datatype):
             # Allow int signals to represent float records when prec is 0
@@ -348,6 +370,8 @@ def make_converter(datatype: Optional[Type], values: Dict[str, Any]) -> PvaConve
                 raise TypeError(f"{pv} has type {typ.__name__} not {datatype.__name__}")
         return PvaConverter()
     elif "NTTable" in typeid:
+        if is_pydantic_model(datatype):
+            return PvaPydanticModelConverter(datatype)  # type: ignore
         return PvaTableConverter()
     elif "structure" in typeid:
         return PvaDictConverter()
@@ -356,15 +380,40 @@ def make_converter(datatype: Optional[Type], values: Dict[str, Any]) -> PvaConve
 
 
 class PvaSignalBackend(SignalBackend[T]):
-    _ctxt: Optional[Context] = None
+    _ctxt: Context | None = None
+
+    _ALLOWED_DATATYPES = (
+        bool,
+        int,
+        float,
+        str,
+        Sequence,
+        np.ndarray,
+        Enum,
+        RuntimeSubsetEnum,
+        BaseModel,
+        dict,
+    )
 
-    def __init__(self, datatype: Optional[Type[T]], read_pv: str, write_pv: str):
+    @classmethod
+    def datatype_allowed(cls, dtype: Any) -> bool:
+        stripped_origin = get_origin(dtype) or dtype
+        if dtype is None:
+            return True
+        return inspect.isclass(stripped_origin) and issubclass(
+            stripped_origin, cls._ALLOWED_DATATYPES
+        )
+
+    def __init__(self, datatype: type[T] | None, read_pv: str, write_pv: str):
         self.datatype = datatype
+        if not PvaSignalBackend.datatype_allowed(self.datatype):
+            raise TypeError(f"Given datatype {self.datatype} unsupported in PVA.")
+
         self.read_pv = read_pv
         self.write_pv = write_pv
-        self.initial_values: Dict[str, Any] = {}
+        self.initial_values: dict[str, Any] = {}
         self.converter: PvaConverter = DisconnectedPvaConverter()
-        self.subscription: Optional[Subscription] = None
+        self.subscription: Subscription | None = None
 
     def source(self, name: str):
         return f"pva://{self.read_pv}"
@@ -404,7 +453,7 @@ class PvaSignalBackend(SignalBackend[T]):
             await self._store_initial_value(self.read_pv, timeout=timeout)
         self.converter = make_converter(self.datatype, self.initial_values)
 
-    async def put(self, value: Optional[T], wait=True, timeout=None):
+    async def put(self, value: T | None, wait=True, timeout=None):
         if value is None:
             write_value = self.initial_values[self.write_pv]
         else:
@@ -424,7 +473,7 @@ class PvaSignalBackend(SignalBackend[T]):
         value = await self.ctxt.get(self.read_pv)
         return self.converter.get_datakey(source, value)
 
-    def _pva_request_string(self, fields: List[str]) -> str:
+    def _pva_request_string(self, fields: list[str]) -> str:
         """
         Converts a list of requested fields into a PVA request string which can be
         passed to p4p.
@@ -447,7 +496,7 @@ class PvaSignalBackend(SignalBackend[T]):
         value = await self.ctxt.get(self.write_pv, "field(value)")
         return self.converter.value(value)
 
-    def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:
+    def set_callback(self, callback: ReadingValueCallback[T] | None) -> None:
         if callback:
             assert (
                 not self.subscription

ophyd_async/epics/signal/_signal.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-from typing import Optional, Tuple, Type
-
 from ophyd_async.core import (
     SignalBackend,
     SignalR,
@@ -19,7 +17,7 @@ from ._epics_transport import _EpicsTransport
 _default_epics_transport = _EpicsTransport.ca
 
 
-def _transport_pv(pv: str) -> Tuple[_EpicsTransport, str]:
+def _transport_pv(pv: str) -> tuple[_EpicsTransport, str]:
     split = pv.split("://", 1)
     if len(split) > 1:
         # We got something like pva://mydevice, so use specified comms mode
@@ -32,7 +30,7 @@ def _transport_pv(pv: str) -> Tuple[_EpicsTransport, str]:
 
 
 def _epics_signal_backend(
-    datatype: Optional[Type[T]], read_pv: str, write_pv: str
+    datatype: type[T] | None, read_pv: str, write_pv: str
 ) -> SignalBackend[T]:
     """Create an epics signal backend."""
     r_transport, r_pv = _transport_pv(read_pv)
@@ -42,7 +40,7 @@ def _epics_signal_backend(
 
 
 def epics_signal_rw(
-    datatype: Type[T], read_pv: str, write_pv: Optional[str] = None, name: str = ""
+    datatype: type[T], read_pv: str, write_pv: str | None = None, name: str = ""
 ) -> SignalRW[T]:
     """Create a `SignalRW` backed by 1 or 2 EPICS PVs
 
@@ -60,7 +58,7 @@ def epics_signal_rw(
 
 
 def epics_signal_rw_rbv(
-    datatype: Type[T], write_pv: str, read_suffix: str = "_RBV", name: str = ""
+    datatype: type[T], write_pv: str, read_suffix: str = "_RBV", name: str = ""
 ) -> SignalRW[T]:
     """Create a `SignalRW` backed by 1 or 2 EPICS PVs, with a suffix on the readback pv
 
@@ -76,7 +74,7 @@ def epics_signal_rw_rbv(
     return epics_signal_rw(datatype, f"{write_pv}{read_suffix}", write_pv, name)
 
 
-def epics_signal_r(datatype: Type[T], read_pv: str, name: str = "") -> SignalR[T]:
+def epics_signal_r(datatype: type[T], read_pv: str, name: str = "") -> SignalR[T]:
     """Create a `SignalR` backed by 1 EPICS PV
 
     Parameters
@@ -90,7 +88,7 @@ def epics_signal_r(datatype: Type[T], read_pv: str, name: str = "") -> SignalR[T
     return SignalR(backend, name=name)
 
 
-def epics_signal_w(datatype: Type[T], write_pv: str, name: str = "") -> SignalW[T]:
+def epics_signal_w(datatype: type[T], write_pv: str, name: str = "") -> SignalW[T]:
     """Create a `SignalW` backed by 1 EPICS PVs
 
     Parameters

ophyd_async/fastcs/panda/__init__.py

@@ -15,10 +15,7 @@ from ._table import (
     DatasetTable,
     PandaHdf5DatasetType,
     SeqTable,
-    SeqTableRow,
     SeqTrigger,
-    seq_table_from_arrays,
-    seq_table_from_rows,
 )
 from ._trigger import (
     PcompInfo,
@@ -45,10 +42,7 @@ __all__ = [
     "DatasetTable",
     "PandaHdf5DatasetType",
     "SeqTable",
-    "SeqTableRow",
     "SeqTrigger",
-    "seq_table_from_arrays",
-    "seq_table_from_rows",
     "PcompInfo",
     "SeqTableInfo",
     "StaticPcompTriggerLogic",

ophyd_async/fastcs/panda/_block.py

@@ -7,6 +7,12 @@ from ophyd_async.core import Device, DeviceVector, SignalR, SignalRW, SubsetEnum
 from ._table import DatasetTable, SeqTable
 
 
+class CaptureMode(str, Enum):
+    FIRST_N = "FIRST_N"
+    LAST_N = "LAST_N"
+    FOREVER = "FOREVER"
+
+
 class DataBlock(Device):
     # In future we may decide to make hdf_* optional
     hdf_directory: SignalRW[str]
@@ -15,6 +21,7 @@ class DataBlock(Device):
     num_captured: SignalR[int]
     create_directory: SignalRW[int]
     directory_exists: SignalR[bool]
+    capture_mode: SignalRW[CaptureMode]
     capture: SignalRW[bool]
     flush_period: SignalRW[float]
     datasets: SignalR[DatasetTable]