wip: pushing because my computer broke

src/ophyd_async/core/_table.py

@@ -1,42 +1,107 @@
+from enum import Enum
+from typing import get_args
+
 import numpy as np
 from pydantic import BaseModel, ConfigDict, model_validator
 
 
+def _concat(value1, value2):
+ if isinstance(value1, np.ndarray):
+ return np.concatenate((value1, value2))
+ else:
+ return value1 + value2
+
+
 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),
+ arrayified_kwargs = {}
+ for field_name, field_value in sub_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]
+ 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 {sub_cls.__name__}."
  )
- 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}."
- )
+ if not isinstance(right, type(self)):
+ raise RuntimeError(
+ 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))
+ field_name: _concat(
+ getattr(self, field_name), getattr(right, field_name)
  )
  for field_name in self.model_fields
  }
  )
 
+ def numpy_dtype(self) -> np.dtype:
+ dtype = []
+ for field_value in self.model_fields.values():
+ if isinstance(field_value, np.ndarray):
+ dtype.append(field_value.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(np.dtype(f"<U{max_length_in_enum}"))
+
+ def numpy_table(self):
+ return np.array(
+ self.numpy_columns(),
+ dtype=self.numpy_dtype(),
+ ).transpose()
+
+ 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_value in self.model_fields.values():
+ if isinstance(field_value, np.ndarray):
+ columns.append(field_value)
+ else:
+ enum_type = get_args(field_value.field_info.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_values, dtype=dtype))
+
+ return columns
+
  @model_validator(mode="after")
  def validate_arrays(self) -> "Table":
  first_length = len(next(iter(self))[1])
@@ -45,10 +110,15 @@ def validate_arrays(self) -> "Table":
  ), "Rows should all be of equal size."
 
  if not all(
- np.issubdtype(
- self.model_fields[field_name].default_factory().dtype, field_value.dtype
+ # 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(),
+ np.ndarray,
  )
- for field_name, field_value in self
+ 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__}`, "

src/ophyd_async/fastcs/panda/_table.py

@@ -51,13 +51,7 @@ class SeqTrigger(str, Enum):
  ),
  Field(default_factory=lambda: np.array([], dtype=np.bool_)),
 ]
-TriggerStr = Annotated[
- np.ndarray[tuple[int], np.unicode_],
- NpArrayPydanticAnnotation.factory(
- data_type=np.unicode_, dimensions=1, strict_data_typing=False
- ),
- Field(default_factory=lambda: np.array([], dtype=np.dtype("<U32"))),
-]
+TriggerStr = Annotated[list[SeqTrigger], Field(default_factory=list)]
 
 
 class SeqTable(Table):
@@ -103,40 +97,12 @@ def row(
  ) -> "SeqTable":
  sig = inspect.signature(cls.row)
  kwargs = {k: v for k, v in locals().items() if k in sig.parameters}
-
- if isinstance(kwargs["trigger"], SeqTrigger):
- kwargs["trigger"] = kwargs["trigger"].value
- elif isinstance(kwargs["trigger"], str):
- SeqTrigger(kwargs["trigger"])
-
+ if not isinstance(kwargs["trigger"], SeqTrigger):
+ if kwargs["trigger"] not in SeqTrigger.__members__.values():
+ raise ValueError(f"'{kwargs['trigger']}' is not a valid trigger.")
+ kwargs["trigger"] = SeqTrigger(kwargs["trigger"])
  return Table.row(cls, **kwargs)
 
- @field_validator("trigger", mode="before")
- @classmethod
- def trigger_to_np_array(cls, trigger_column):
- """
- The user can provide a list of SeqTrigger enum elements instead of a numpy str.
- """
- if isinstance(trigger_column, Sequence) and all(
- isinstance(trigger, SeqTrigger) for trigger in trigger_column
- ):
- trigger_column = np.array(
- [trigger.value for trigger in trigger_column], dtype=np.dtype("<U32")
- )
- elif isinstance(trigger_column, Sequence) or isinstance(
- trigger_column, np.ndarray
- ):
- for trigger in trigger_column:
- SeqTrigger(
- trigger
- ) # To check all the given strings are actually `SeqTrigger`s
- else:
- raise ValueError(
- "Expected a numpy array or a sequence of `SeqTrigger`, got "
- f"{type(trigger_column)}."
- )
- return trigger_column
-
  @model_validator(mode="after")
  def validate_max_length(self) -> "SeqTable":
  """

tests/fastcs/panda/test_table.py

@@ -12,8 +12,11 @@ def test_seq_table_converts_lists():
  seq_table_dict_with_lists = {field_name: [] for field_name, _ in SeqTable()}
  # Validation passes
  seq_table = SeqTable(**seq_table_dict_with_lists)
- assert isinstance(seq_table.trigger, np.ndarray)
- assert seq_table.trigger.dtype == np.dtype("U32")
+ for field_name, field_value in seq_table:
+ if field_name == "trigger":
+ assert field_value == []
+ else:
+ assert np.array_equal(field_value, np.array([], dtype=field_value.dtype))
 
 
 def test_seq_table_validation_errors():
@@ -73,16 +76,20 @@ def test_seq_table_validation_errors():
  wrong_types = {
  field_name: field_value.astype(np.unicode_)
  for field_name, field_value in row_one
+ if isinstance(field_value, np.ndarray)
  }
  SeqTable(**wrong_types)
 
 
 def test_seq_table_pva_conversion():
  pva_dict = {
  "repeats": np.array([1, 2, 3, 4], dtype=np.int32),
- "trigger": np.array(
- ["Immediate", "Immediate", "BITC=0", "Immediate"], dtype=np.dtype("U32")
- ),
+ "trigger": [
+ SeqTrigger.IMMEDIATE,
+ SeqTrigger.IMMEDIATE,
+ SeqTrigger.BITC_0,
+ SeqTrigger.IMMEDIATE,
+ ],
  "position": np.array([1, 2, 3, 4], dtype=np.int32),
  "time1": np.array([1, 0, 1, 0], dtype=np.int32),
  "outa1": np.array([1, 0, 1, 0], dtype=np.bool_),
@@ -178,31 +185,36 @@ def test_seq_table_pva_conversion():
  },
  ]
 
+ def _assert_col_equal(column1, column2):
+ if isinstance(column1, np.ndarray):
+ assert np.array_equal(column1, column2)
+ assert column1.dtype == column2.dtype
+ else:
+ assert column1 == column2
+ assert all(isinstance(x, SeqTrigger) for x in column1)
+ assert all(isinstance(x, SeqTrigger) for x in column2)
+
  seq_table_from_pva_dict = SeqTable(**pva_dict)
  for (_, column1), column2 in zip(seq_table_from_pva_dict, pva_dict.values()):
- assert np.array_equal(column1, column2)
- assert column1.dtype == column2.dtype
+ _assert_col_equal(column1, column2)
 
  seq_table_from_rows = reduce(
  lambda x, y: x + y,
  [SeqTable.row(**row_kwargs) for row_kwargs in row_wise_dicts],
  )
  for (_, column1), column2 in zip(seq_table_from_rows, pva_dict.values()):
- assert np.array_equal(column1, column2)
- assert column1.dtype == column2.dtype
+ _assert_col_equal(column1, column2)
 
  # Idempotency
  applied_twice_to_pva_dict = SeqTable(**pva_dict).model_dump(mode="python")
  for column1, column2 in zip(applied_twice_to_pva_dict.values(), pva_dict.values()):
- assert np.array_equal(column1, column2)
- assert column1.dtype == column2.dtype
+ _assert_col_equal(column1, column2)
 
 
 def test_seq_table_takes_trigger_enum_row():
  for trigger in (SeqTrigger.BITA_0, "BITA=0"):
  table = SeqTable.row(trigger=trigger)
- assert table.trigger[0] == "BITA=0"
- assert np.issubdtype(table.trigger.dtype, np.dtype("<U32"))
+ assert table.trigger[0] == SeqTrigger.BITA_0
  table = SeqTable(
  repeats=np.array([1], dtype=np.int32),
  trigger=[trigger],
@@ -222,5 +234,4 @@ def test_seq_table_takes_trigger_enum_row():
  oute2=np.array([1], dtype=np.bool_),
  outf2=np.array([1], dtype=np.bool_),
  )
- assert table.trigger[0] == "BITA=0"
- assert np.issubdtype(table.trigger.dtype, np.dtype("<U32"))
+ assert table.trigger[0] == SeqTrigger.BITA_0

tests/fastcs/panda/test_trigger.py

@@ -144,5 +144,5 @@ def full_seq_table(trigger):
  full_seq_table(["A"])
  assert "Value error, 'A' is not a valid SeqTrigger" in str(exc)
  with pytest.raises(ValidationError) as exc:
- full_seq_table({"A"})
+ full_seq_table({"Immediate"})
  assert "Expected a numpy array or a sequence of `SeqTrigger`, got" in str(exc)