feat: make Axis public (#201)

Author: tlambert03

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",

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
 

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,