feat: new iteration pattern/rewrite (take 2)

pyproject.toml

@@ -37,7 +37,7 @@ classifiers = [
 ]
 dynamic = ["version"]
 dependencies = [
-    "pydantic >=2.6",
+    "pydantic >=2.10",
     "numpy >=2.1.0; python_version >= '3.13'",
     "numpy >=1.26.0; python_version >= '3.12'",
     "numpy >=1.25.2",
@@ -123,15 +123,14 @@ keep-runtime-typing = true
 [tool.ruff.lint.per-file-ignores]
 "tests/*.py" = ["D", "S101", "E501", "SLF"]
 
-[tool.ruff.lint.flake8-tidy-imports]
-# Disallow all relative imports.
-ban-relative-imports = "all"
-
 # https://docs.pytest.org/en/6.2.x/customize.html
 [tool.pytest.ini_options]
 minversion = "6.0"
 testpaths = ["tests"]
-filterwarnings = ["error"]
+filterwarnings = [
+    "error",
+    "ignore:.*Positions no longer have a sequence attribute",
+]
 
 # https://mypy.readthedocs.io/en/stable/config_file.html
 [tool.mypy]
@@ -143,6 +142,10 @@ show_error_codes = true
 pretty = true
 plugins = ["pydantic.mypy"]
 
+[tool.pyright]
+include = ["src", "tests/v2"]
+reportArgumentType = false
+
 # https://coverage.readthedocs.io/en/6.4/config.html
 [tool.coverage.run]
 source = ["useq"]

src/useq/__init__.py

@@ -5,18 +5,18 @@
 
 from useq._actions import AcquireImage, Action, CustomAction, HardwareAutofocus
 from useq._channel import Channel
+from useq._enums import Axis, RelativeTo, Shape
 from useq._grid import (
     GridFromEdges,
     GridRowsColumns,
     GridWidthHeight,
     MultiPointPlan,
     RandomPoints,
     RelativeMultiPointPlan,
-    Shape,
 )
 from useq._hardware_autofocus import AnyAutofocusPlan, AutoFocusPlan, AxesBasedAF
 from useq._mda_event import Channel as EventChannel
-from useq._mda_event import MDAEvent, PropertyTuple, SLMImage
+from useq._mda_event import MDAEvent, MutableMDAEvent, PropertyTuple, SLMImage
 from useq._mda_sequence import MDASequence
 from useq._plate import WellPlate, WellPlatePlan
 from useq._plate_registry import register_well_plates, registered_well_plate_keys
@@ -29,7 +29,6 @@
     TIntervalDuration,
     TIntervalLoops,
 )
-from useq._utils import Axis
 from useq._z import (
     AnyZPlan,
     ZAboveBelow,
@@ -65,12 +64,14 @@
     "MDASequence",
     "MultiPhaseTimePlan",
     "MultiPointPlan",
+    "MutableMDAEvent",
     "OrderMode",
     "Position",  # alias for AbsolutePosition
     "PropertyTuple",
     "RandomPoints",
     "RelativeMultiPointPlan",
     "RelativePosition",
+    "RelativeTo",
     "SLMImage",
     "Shape",
     "TDurationLoops",

src/useq/_base_model.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from re import findall
 from types import MappingProxyType
 from typing import (
     IO,
@@ -27,10 +26,11 @@
 _T = TypeVar("_T", bound="FrozenModel")
 _Y = TypeVar("_Y", bound="UseqModel")
 
-PYDANTIC_VERSION = tuple(int(x) for x in findall(r"\d_", pydantic.__version__)[:3])
-GET_DEFAULT_KWARGS: dict = {}
+PYDANTIC_VERSION = tuple(int(x) for x in pydantic.__version__.split(".")[:2])
 if PYDANTIC_VERSION >= (2, 10):
-    GET_DEFAULT_KWARGS = {"validated_data": {}}
+    GET_DEFAULT_KWARGS: dict = {"validated_data": {}}
+else:
+    GET_DEFAULT_KWARGS = {}
 
 
 class _ReplaceableModel(BaseModel):
@@ -87,9 +87,9 @@
     )
 
 
-class UseqModel(FrozenModel):
+class IOMixin(BaseModel):
     @classmethod
-    def from_file(cls: type[_Y], path: Union[str, Path]) -> _Y:
+    def from_file(cls, path: Union[str, Path]) -> "Self":
         """Return an instance of this class from a file.  Supports JSON and YAML."""
         path = Path(path)
         if path.suffix in {".yaml", ".yml"}:
@@ -148,3 +148,9 @@
             exclude_none=exclude_none,
         )
         return yaml.safe_dump(data, stream=stream)
+
+
+class MutableUseqModel(IOMixin, MutableModel): ...
+
+
+class UseqModel(FrozenModel, IOMixin): ...

src/useq/_enums.py

@@ -0,0 +1,69 @@
+from enum import Enum
+from typing import Final, Literal
+
+
+class Axis(str, Enum):
+    """Recognized useq-schema axis keys.
+
+    Attributes
+    ----------
+    TIME : Literal["t"]
+        Time axis.
+    POSITION : Literal["p"]
+        XY Stage Position axis.
+    GRID : Literal["g"]
+        Grid axis (usually an additional row/column iteration around a position).
+    CHANNEL : Literal["c"]
+        Channel axis.
+    Z : Literal["z"]
+        Z axis.
+    """
+
+    TIME = "t"
+    POSITION = "p"
+    GRID = "g"
+    CHANNEL = "c"
+    Z = "z"
+
+    def __str__(self) -> Literal["t", "p", "g", "c", "z"]:
+        return self.value
+
+
+# note: order affects the default axis_order in MDASequence
+AXES: Final[tuple[Axis, ...]] = (
+    Axis.TIME,
+    Axis.POSITION,
+    Axis.GRID,
+    Axis.CHANNEL,
+    Axis.Z,
+)
+
+
+class RelativeTo(Enum):
+    """Where the coordinates of the grid are relative to.
+
+    Attributes
+    ----------
+    center : Literal['center']
+        Grid is centered around the origin.
+    top_left : Literal['top_left']
+        Grid is positioned such that the top left corner is at the origin.
+    """
+
+    center = "center"
+    top_left = "top_left"
+
+
+class Shape(Enum):
+    """Shape of the bounding box for random points.
+
+    Attributes
+    ----------
+    ELLIPSE : Literal['ellipse']
+        The bounding box is an ellipse.
+    RECTANGLE : Literal['rectangle']
+        The bounding box is a rectangle.
+    """
+
+    ELLIPSE = "ellipse"
+    RECTANGLE = "rectangle"

src/useq/_grid.py

@@ -4,21 +4,22 @@
 import math
 import warnings
 from collections.abc import Iterable, Iterator, Sequence
-from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Annotated,
     Any,
     Callable,
+    Generic,
     Optional,
     Union,
 )
 
 import numpy as np
 from annotated_types import Ge, Gt
-from pydantic import Field, field_validator, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self, TypeAlias
 
+from useq._enums import RelativeTo, Shape
 from useq._point_visiting import OrderMode, TraversalOrder
 from useq._position import (
     AbsolutePosition,
@@ -37,47 +38,15 @@
 MIN_RANDOM_POINTS = 10000
 
 
-class RelativeTo(Enum):
-    """Where the coordinates of the grid are relative to.
-
-    Attributes
-    ----------
-    center : Literal['center']
-        Grid is centered around the origin.
-    top_left : Literal['top_left']
-        Grid is positioned such that the top left corner is at the origin.
-    """
-
-    center = "center"
-    top_left = "top_left"
-
-
-# used in iter_indices below, to determine the order in which indices are yielded
-class _GridPlan(_MultiPointPlan[PositionT]):
-    """Base class for all grid plans.
-
-    Attributes
-    ----------
-    overlap : float | Tuple[float, float]
-        Overlap between grid positions in percent. If a single value is provided, it is
-        used for both x and y. If a tuple is provided, the first value is used
-        for x and the second for y.
-    mode : OrderMode
-        Define the ways of ordering the grid positions. Options are
-        row_wise, column_wise, row_wise_snake, column_wise_snake and spiral.
-        By default, row_wise_snake.
-    fov_width : Optional[float]
-        Width of the field of view in microns.  If not provided, acquisition engines
-        should use current width of the FOV based on the current objective and camera.
-        Engines MAY override this even if provided.
-    fov_height : Optional[float]
-        Height of the field of view in microns. If not provided, acquisition engines
-        should use current height of the FOV based on the current objective and camera.
-        Engines MAY override this even if provided.
-    """
-
+class _GridMixin(BaseModel, Generic[PositionT]):
     overlap: tuple[float, float] = Field(default=(0.0, 0.0), frozen=True)
     mode: OrderMode = Field(default=OrderMode.row_wise_snake, frozen=True)
+    fov_width: Optional[float] = None
+    fov_height: Optional[float] = None
+
+    @property
+    def is_relative(self) -> bool:
+        return True
 
     @field_validator("overlap", mode="before")
     def _validate_overlap(cls, v: Any) -> tuple[float, float]:
@@ -105,24 +74,18 @@ def _ncolumns(self, dx: float) -> int:
         """Return the number of columns, given a grid step size."""
         raise NotImplementedError
 
-    def num_positions(self) -> int:
-        """Return the number of individual positions in the grid.
+    def __iter__(self) -> Iterator[PositionT]:  # type: ignore [override]
+        yield from self.iter_grid_positions()
 
-        Note: For GridFromEdges and GridWidthHeight, this will depend on field of view
-        size. If no field of view size is provided, the number of positions will be 1.
-        """
-        if isinstance(self, (GridFromEdges, GridWidthHeight)) and (
-            self.fov_width is None or self.fov_height is None
-        ):
-            raise ValueError(
-                "Retrieving the number of positions in a GridFromEdges or "
-                "GridWidthHeight plan requires the field of view size to be set."
-            )
+    def _step_size(self, fov_width: float, fov_height: float) -> tuple[float, float]:
+        dx = fov_width - (fov_width * self.overlap[0]) / 100
+        dy = fov_height - (fov_height * self.overlap[1]) / 100
+        return dx, dy
 
-        dx, dy = self._step_size(self.fov_width or 1, self.fov_height or 1)
-        rows = self._nrows(dy)
-        cols = self._ncolumns(dx)
-        return rows * cols
+    def _build_position(self, **kwargs: Any) -> PositionT:
+        """Build a position object for this grid plan."""
+        pos_cls = RelativePosition if self.is_relative else AbsolutePosition
+        return pos_cls(**kwargs)  # type: ignore
 
     def iter_grid_positions(
         self,
@@ -142,26 +105,36 @@ def iter_grid_positions(
         x0 = self._offset_x(dx)
         y0 = self._offset_y(dy)
 
-        pos_cls = RelativePosition if self.is_relative else AbsolutePosition
         for idx, (r, c) in enumerate(order.generate_indices(rows, cols)):
-            yield pos_cls(  # type: ignore [misc]
+            yield self._build_position(
                 x=x0 + c * dx,
                 y=y0 - r * dy,
                 row=r,
                 col=c,
                 name=f"{str(idx).zfill(4)}",
             )
 
-    def __iter__(self) -> Iterator[PositionT]:  # type: ignore [override]
-        yield from self.iter_grid_positions()
+    def num_positions(self) -> int:
+        """Return the number of individual positions in the grid.
 
-    def _step_size(self, fov_width: float, fov_height: float) -> tuple[float, float]:
-        dx = fov_width - (fov_width * self.overlap[0]) / 100
-        dy = fov_height - (fov_height * self.overlap[1]) / 100
-        return dx, dy
+        Note: For GridFromEdges and GridWidthHeight, this will depend on field of view
+        size. If no field of view size is provided, the number of positions will be 1.
+        """
+        if isinstance(self, (GridFromEdges, GridWidthHeight)) and (
+            self.fov_width is None or self.fov_height is None
+        ):
+            raise ValueError(
+                "Retrieving the number of positions in a GridFromEdges or "
+                "GridWidthHeight plan requires the field of view size to be set."
+            )
+
+        dx, dy = self._step_size(self.fov_width or 1, self.fov_height or 1)
+        rows = self._nrows(dy)
+        cols = self._ncolumns(dx)
+        return rows * cols
 
 
-class GridFromEdges(_GridPlan[AbsolutePosition]):
+class GridFromEdges(_GridMixin, _MultiPointPlan[AbsolutePosition]):
     """Yield absolute stage positions to cover a bounded area.
 
     The bounded area is defined by top, left, bottom and right edges in
@@ -253,7 +226,7 @@ def plot(self, *, show: bool = True) -> Axes:
         )
 
 
-class GridRowsColumns(_GridPlan[RelativePosition]):
+class GridRowsColumns(_GridMixin, _MultiPointPlan[RelativePosition]):
     """Grid plan based on number of rows and columns.
 
     Attributes
@@ -311,7 +284,7 @@ def _offset_y(self, dy: float) -> float:
 GridRelative = GridRowsColumns
 
 
-class GridWidthHeight(_GridPlan[RelativePosition]):
+class GridWidthHeight(_GridMixin, _MultiPointPlan[RelativePosition]):
     """Grid plan based on total width and height.
 
     Attributes
@@ -371,21 +344,6 @@ def _offset_y(self, dy: float) -> float:
 # ------------------------ RANDOM ------------------------
 
 
-class Shape(Enum):
-    """Shape of the bounding box for random points.
-
-    Attributes
-    ----------
-    ELLIPSE : Literal['ellipse']
-        The bounding box is an ellipse.
-    RECTANGLE : Literal['rectangle']
-        The bounding box is a rectangle.
-    """
-
-    ELLIPSE = "ellipse"
-    RECTANGLE = "rectangle"
-
-
 class RandomPoints(_MultiPointPlan[RelativePosition]):
     """Yield random points in a specified geometric shape.