Merge branch 'main' into 549-support-multi-run-in-standarddetectorpre…

…pares

src/ophyd_async/core/__init__.py

@@ -61,9 +61,14 @@
  soft_signal_rw,
  wait_for_value,
 )
-from ._signal_backend import RuntimeSubsetEnum, SignalBackend, SubsetEnum
+from ._signal_backend import (
+ RuntimeSubsetEnum,
+ SignalBackend,
+ SubsetEnum,
+)
 from ._soft_signal_backend import SignalMetadata, SoftSignalBackend
 from ._status import AsyncStatus, WatchableAsyncStatus, completed_status
+from ._table import Table
 from ._utils import (
  DEFAULT_TIMEOUT,
  CalculatableTimeout,
@@ -152,6 +157,7 @@
  "CalculateTimeout",
  "NotConnected",
  "ReadingValueCallback",
+ "Table",
  "T",
  "WatcherUpdate",
  "get_dtype",

src/ophyd_async/core/_device_save_loader.py

@@ -7,6 +7,7 @@
 from bluesky.plan_stubs import abs_set, wait
 from bluesky.protocols import Location
 from bluesky.utils import Msg
+from pydantic import BaseModel
 
 from ._device import Device
 from ._signal import SignalRW
@@ -18,6 +19,12 @@ def ndarray_representer(dumper: yaml.Dumper, array: npt.NDArray[Any]) -> yaml.No
  )
 
 
+def pydantic_model_abstraction_representer(
+ dumper: yaml.Dumper, model: BaseModel
+) -> yaml.Node:
+ return dumper.represent_data(model.model_dump(mode="python"))
+
+
 class OphydDumper(yaml.Dumper):
  def represent_data(self, data: Any) -> Any:
  if isinstance(data, Enum):
@@ -134,6 +141,11 @@ def save_to_yaml(phases: Sequence[Dict[str, Any]], save_path: str) -> None:
  """
 
  yaml.add_representer(np.ndarray, ndarray_representer, Dumper=yaml.Dumper)
+ yaml.add_multi_representer(
+ BaseModel,
+ pydantic_model_abstraction_representer,
+ Dumper=yaml.Dumper,
+ )
 
  with open(save_path, "w") as file:
  yaml.dump(phases, file, Dumper=OphydDumper, default_flow_style=False)

src/ophyd_async/core/_signal_backend.py

@@ -1,5 +1,13 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, ClassVar, Generic, Literal, Optional, Tuple, Type
+from typing import (
+ TYPE_CHECKING,
+ ClassVar,
+ Generic,
+ Literal,
+ Optional,
+ Tuple,
+ Type,
+)
 
 from ._protocol import DataKey, Reading
 from ._utils import DEFAULT_TIMEOUT, ReadingValueCallback, T
@@ -11,6 +19,11 @@ class SignalBackend(Generic[T]):
  #: Datatype of the signal value
  datatype: Optional[Type[T]] = None
 
+ @classmethod
+ @abstractmethod
+ def datatype_allowed(cls, dtype: type):
+ """Check if a given datatype is acceptable for this signal backend."""
+
  #: Like ca://PV_PREFIX:SIGNAL
  @abstractmethod
  def source(self, name: str) -> str:

src/ophyd_async/core/_soft_signal_backend.py

@@ -2,15 +2,20 @@
 
 import inspect
 import time
+from abc import ABCMeta
 from collections import abc
 from enum import Enum
 from typing import Dict, Generic, Optional, Tuple, Type, Union, cast, get_origin
 
 import numpy as np
 from bluesky.protocols import DataKey, Dtype, Reading
+from pydantic import BaseModel
 from typing_extensions import TypedDict
 
-from ._signal_backend import RuntimeSubsetEnum, SignalBackend
+from ._signal_backend import (
+ RuntimeSubsetEnum,
+ SignalBackend,
+)
 from ._utils import DEFAULT_TIMEOUT, ReadingValueCallback, T, get_dtype
 
 primitive_dtypes: Dict[type, Dtype] = {
@@ -94,7 +99,7 @@ def make_initial_value(self, datatype: Optional[Type[T]]) -> T:
 class SoftEnumConverter(SoftConverter):
  choices: Tuple[str, ...]
 
- def __init__(self, datatype: Union[RuntimeSubsetEnum, Enum]):
+ def __init__(self, datatype: Union[RuntimeSubsetEnum, Type[Enum]]):
  if issubclass(datatype, Enum):
  self.choices = tuple(v.value for v in datatype)
  else:
@@ -122,17 +127,36 @@ def make_initial_value(self, datatype: Optional[Type[T]]) -> T:
  return cast(T, self.choices[0])
 
 
+class SoftPydanticModelConverter(SoftConverter):
+ def __init__(self, datatype: Type[BaseModel]):
+ self.datatype = datatype
+
+ def write_value(self, value):
+ if isinstance(value, dict):
+ return self.datatype(**value)
+ return value
+
+
 def make_converter(datatype):
  is_array = get_dtype(datatype) is not None
  is_sequence = get_origin(datatype) == abc.Sequence
  is_enum = inspect.isclass(datatype) and (
  issubclass(datatype, Enum) or issubclass(datatype, RuntimeSubsetEnum)
  )
 
+ is_pydantic_model = (
+ inspect.isclass(datatype)
+ # Necessary to avoid weirdness in ABCMeta.__subclasscheck__
+ and isinstance(datatype, ABCMeta)
+ and issubclass(datatype, BaseModel)
+ )
+
  if is_array or is_sequence:
  return SoftArrayConverter()
  if is_enum:
  return SoftEnumConverter(datatype)
+ if is_pydantic_model:
+ return SoftPydanticModelConverter(datatype)
 
  return SoftConverter()
 
@@ -145,6 +169,10 @@ class SoftSignalBackend(SignalBackend[T]):
  _timestamp: float
  _severity: int
 
+ @classmethod
+ def datatype_allowed(cls, datatype: Type) -> bool:
+ return True # Any value allowed in a soft signal
+
  def __init__(
  self,
  datatype: Optional[Type[T]],

src/ophyd_async/core/_table.py

@@ -0,0 +1,58 @@
+import numpy as np
+from pydantic import BaseModel, ConfigDict, model_validator
+
+
+class Table(BaseModel):
+ """An abstraction of a Table of str to numpy array."""
+
+ model_config = ConfigDict(validate_assignment=True, strict=False)
+
+ @classmethod
+ def row(cls, sub_cls, **kwargs) -> "Table":
+ arrayified_kwargs = {
+ field_name: np.concatenate(
+ (
+ (default_arr := field_value.default_factory()),
+ np.array([kwargs[field_name]], dtype=default_arr.dtype),
+ )
+ )
+ for field_name, field_value in sub_cls.model_fields.items()
+ }
+ return sub_cls(**arrayified_kwargs)
+
+ def __add__(self, right: "Table") -> "Table":
+ """Concatenate the arrays in field values."""
+
+ assert isinstance(right, type(self)), (
+ f"{right} is not a `Table`, or is not the same "
+ f"type of `Table` as {self}."
+ )
+
+ return type(self)(
+ **{
+ field_name: np.concatenate(
+ (getattr(self, field_name), getattr(right, field_name))
+ )
+ for field_name in self.model_fields
+ }
+ )
+
+ @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(
+ np.issubdtype(
+ self.model_fields[field_name].default_factory().dtype, field_value.dtype
+ )
+ for field_name, field_value in self
+ ):
+ raise ValueError(
+ f"Cannot construct a `{type(self).__name__}`, "
+ "some rows have incorrect types."
+ )
+
+ return self

src/ophyd_async/core/_utils.py

@@ -145,7 +145,7 @@ def get_dtype(typ: Type) -> Optional[np.dtype]:
 
 
 def get_unique(values: Dict[str, T], types: str) -> T:
- """If all values are the same, return that value, otherwise return TypeError
+ """If all values are the same, return that value, otherwise raise TypeError
 
  >>> get_unique({"a": 1, "b": 1}, "integers")
  1

src/ophyd_async/epics/signal/_aioca.py

@@ -1,9 +1,10 @@
+import inspect
 import logging
 import sys
 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, Dict, List, Optional, Sequence, Type, Union, get_origin
 
 import numpy as np
 from aioca import (
@@ -24,6 +25,7 @@
  DEFAULT_TIMEOUT,
  NotConnected,
  ReadingValueCallback,
+ RuntimeSubsetEnum,
  SignalBackend,
  T,
  get_dtype,
@@ -211,7 +213,8 @@ def make_converter(
  raise TypeError(f"{pv} has type [{pv_dtype}] not [{dtype}]")
  return CaArrayConverter(pv_dbr, None)
  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 +243,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)
 
 
 _tried_pyepics = False
@@ -256,8 +259,31 @@ def _use_pyepics_context_if_imported():
 
 
 class CaSignalBackend(SignalBackend[T]):
+ _ALLOWED_DATATYPES = (
+ bool,
+ int,
+ float,
+ str,
+ Sequence,
+ Enum,
+ RuntimeSubsetEnum,
+ np.ndarray,
+ )
+
+ @classmethod
+ def datatype_allowed(cls, datatype: Optional[Type]) -> bool:
+ stripped_origin = get_origin(datatype) or datatype
+ if datatype is None:
+ return True
+
+ return inspect.isclass(stripped_origin) and issubclass(
+ stripped_origin, cls._ALLOWED_DATATYPES
+ )
+
  def __init__(self, datatype: Optional[Type[T]], 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] = {}

src/ophyd_async/epics/signal/_p4p.py

@@ -3,14 +3,17 @@
 import inspect
 import logging
 import time
+from abc import ABCMeta
 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, Dict, List, Optional, Sequence, Type, Union, get_origin
 
+import numpy as np
 from bluesky.protocols import DataKey, Dtype, Reading
 from p4p import Value
 from p4p.client.asyncio import Context, Subscription
+from pydantic import BaseModel
 
 from ophyd_async.core import (
  DEFAULT_TIMEOUT,
@@ -253,6 +256,19 @@ def get_datakey(self, source: str, value) -> DataKey:
  return _data_key_from_value(source, value, dtype="object")
 
 
+class PvaPydanticModelConverter(PvaConverter):
+ def __init__(self, datatype: BaseModel):
+ self.datatype = datatype
+
+ def value(self, value: Value):
+ return self.datatype(**value.todict())
+
+ def write_value(self, value: Union[BaseModel, Dict[str, Any]]):
+ if isinstance(value, self.datatype):
+ return value.model_dump(mode="python")
+ return value
+
+
 class PvaDictConverter(PvaConverter):
  def reading(self, value):
  ts = time.time()
@@ -348,6 +364,15 @@ 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 (
+ datatype
+ and inspect.isclass(datatype)
+ and
+ # Necessary to avoid weirdness in ABCMeta.__subclasscheck__
+ isinstance(datatype, ABCMeta)
+ and issubclass(datatype, BaseModel)
+ ):
+ return PvaPydanticModelConverter(datatype)
  return PvaTableConverter()
  elif "structure" in typeid:
  return PvaDictConverter()
@@ -358,8 +383,33 @@ def make_converter(datatype: Optional[Type], values: Dict[str, Any]) -> PvaConve
 class PvaSignalBackend(SignalBackend[T]):
  _ctxt: Optional[Context] = None
 
+ _ALLOWED_DATATYPES = (
+ bool,
+ int,
+ float,
+ str,
+ Sequence,
+ np.ndarray,
+ Enum,
+ RuntimeSubsetEnum,
+ BaseModel,
+ dict,
+ )
+
+ @classmethod
+ def datatype_allowed(cls, datatype: Optional[Type]) -> bool:
+ stripped_origin = get_origin(datatype) or datatype
+ if datatype is None:
+ return True
+ return inspect.isclass(stripped_origin) and issubclass(
+ stripped_origin, cls._ALLOWED_DATATYPES
+ )
+
  def __init__(self, datatype: Optional[Type[T]], 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] = {}