diff --git a/src/useq/__init__.py b/src/useq/__init__.py index 5ac235d..8b0ca64 100644 --- a/src/useq/__init__.py +++ b/src/useq/__init__.py @@ -28,6 +28,7 @@ TIntervalDuration, TIntervalLoops, ) +from useq._utils import Axis from useq._z import ( AnyZPlan, ZAboveBelow, @@ -46,6 +47,7 @@ "AnyZPlan", "AutoFocusPlan", "AxesBasedAF", + "Axis", "Channel", "GridFromEdges", "GridRelative", diff --git a/src/useq/_mda_sequence.py b/src/useq/_mda_sequence.py index ae24b7e..0faaf57 100644 --- a/src/useq/_mda_sequence.py +++ b/src/useq/_mda_sequence.py @@ -345,8 +345,8 @@ def _check_order( ) ): raise ValueError( - f"{Axis.Z!r} cannot precede {Axis.POSITION!r} in acquisition order if " - "any position specifies a z_plan" + f"{str(Axis.Z)!r} cannot precede {str(Axis.POSITION)!r} in acquisition " + "order if any position specifies a z_plan" ) if ( @@ -426,12 +426,12 @@ def used_axes(self) -> str: def iter_axis(self, axis: str) -> Iterator[Channel | float | PositionBase]: """Iterate over the positions or items of a given axis.""" plan = { - Axis.TIME: self.time_plan, - Axis.POSITION: self.stage_positions, - Axis.Z: self.z_plan, - Axis.CHANNEL: self.channels, - Axis.GRID: self.grid_plan, - }[axis] + str(Axis.TIME): self.time_plan, + str(Axis.POSITION): self.stage_positions, + str(Axis.Z): self.z_plan, + str(Axis.CHANNEL): self.channels, + str(Axis.GRID): self.grid_plan, + }[str(axis).lower()] if plan: yield from plan diff --git a/src/useq/_utils.py b/src/useq/_utils.py index 9a0cc87..88fc427 100644 --- a/src/useq/_utils.py +++ b/src/useq/_utils.py @@ -2,10 +2,11 @@ import re from datetime import timedelta -from typing import TYPE_CHECKING, Literal, NamedTuple, TypeVar +from enum import Enum +from typing import TYPE_CHECKING, NamedTuple if TYPE_CHECKING: - from typing import Final + from typing import Final, Literal, TypeVar from typing_extensions import TypeGuard @@ -17,18 +18,35 @@ # could be an enum, but this more easily allows Axis.Z to be a string -class Axis: - """Recognized axis names.""" +class Axis(str, Enum): + """Recognized useq-schema axis keys. - TIME: Final[Literal["t"]] = "t" - POSITION: Final[Literal["p"]] = "p" - GRID: Final[Literal["g"]] = "g" - CHANNEL: Final[Literal["c"]] = "c" - Z: Final[Literal["z"]] = "z" + 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: Literal["t"] = "t" + POSITION: Literal["p"] = "p" + GRID: Literal["g"] = "g" + CHANNEL: Literal["c"] = "c" + Z: Literal["z"] = "z" + + def __str__(self) -> Literal["t", "p", "g", "c", "z"]: + return self.value # type: ignore [no-any-return] # note: order affects the default axis_order in MDASequence -AXES: Final[tuple[str, ...]] = ( +AXES: Final[tuple[Axis, ...]] = ( Axis.TIME, Axis.POSITION, Axis.GRID,