feat: add CustomAction (#214)

Author: tlambert03

src/useq/__init__.py

@@ -3,7 +3,7 @@
 import warnings
 from typing import Any
 
-from useq._actions import AcquireImage, Action, HardwareAutofocus
+from useq._actions import AcquireImage, Action, CustomAction, HardwareAutofocus
 from useq._channel import Channel
 from useq._grid import (
     GridFromEdges,
@@ -50,6 +50,7 @@
     "AxesBasedAF",
     "Axis",
     "Channel",
+    "CustomAction",
     "EventChannel",
     "GridFromEdges",
     "GridRelative",

src/useq/_actions.py

@@ -1,17 +1,24 @@
 from typing import Optional, Union
 
+from pydantic import ConfigDict, Field, TypeAdapter, field_validator
+from pydantic_core import PydanticSerializationError
 from typing_extensions import Literal
 
 from useq._base_model import FrozenModel
 
+_dict_adapter = TypeAdapter(dict, config=ConfigDict(defer_build=True))
+
 
 class Action(FrozenModel):
     """Base class for a [`useq.MDAEvent`][] action.
 
     An `Action` specifies what task should be performed during a
     [`useq.MDAEvent`][]. An `Action` can be for example used to acquire an
     image ([`useq.AcquireImage`][]) or to perform a hardware autofocus
-    ([`useq.HardwareAutofocus`][]).
+    ([`useq.HardwareAutofocus`][]).  An action of `None` implies `AcquireImage`.
+
+    You may use `CustomAction` to indicate any custom action, with the `data` attribute
+    containing any data required to perform the custom action.
 
     Attributes
     ----------
@@ -61,4 +68,42 @@ class HardwareAutofocus(Action):
     max_retries: int = 3
 
 
-AnyAction = Union[HardwareAutofocus, AcquireImage]
+class CustomAction(Action):
+    """[`useq.Action`][] to perform a custom action.
+
+    This is a generic user action that can be used to represent anything that is not
+    covered by the other action types, such as a microfluidic event, or a
+    photostimulation, etc...
+
+    The `data` attribute is a dictionary that can contain any data that is needed to
+    perform the custom action.  It *must* be serializable to JSON by `pydantic`.
+
+    Attributes
+    ----------
+    type : Literal["custom"]
+        This action can be used to perform a custom action.
+    name : str, optional
+        A name for the custom action (not to be confused with the `type` attribute,
+        which must always be `"custom"`).
+    data : dict, optional
+        Custom data associated with the action.
+    """
+
+    type: Literal["custom"] = "custom"
+    name: str = ""
+    data: dict = Field(default_factory=dict)
+
+    @field_validator("data", mode="after")
+    @classmethod
+    def _ensure_serializable(cls, data: dict) -> dict:
+        try:
+            _dict_adapter.serializer.to_json(data)
+        except PydanticSerializationError as e:
+            raise ValueError(
+                f"`CustomAction.data` must be JSON serializable, but is not:\n  {e}.\n"
+                "  (You may use a pydantic object for custom serialization).\n   "
+            ) from e
+        return data
+
+
+AnyAction = Union[HardwareAutofocus, AcquireImage, CustomAction]

src/useq/_mda_event.py

@@ -174,7 +174,10 @@ class MDAEvent(UseqModel):
     action : Action
         The action to perform for this event.  By default, [`useq.AcquireImage`][].
         Example of another action is [`useq.HardwareAutofocus`][] which could be used
-        to perform a hardware autofocus.
+        to perform a hardware autofocus.  For backwards compatibility, an `action` of
+        `None` implies `AcquireImage`.  You may use `CustomAction` to indicate any
+        custom action, with the `data` attribute containing any data required to
+        perform the custom action.
     keep_shutter_open : bool
         If `True`, the illumination shutter should be left open after the event has
         been executed, otherwise it should be closed. By default, `False`."
@@ -199,7 +202,7 @@ class MDAEvent(UseqModel):
     sequence: Optional["MDASequence"] = Field(default=None, repr=False)
     properties: Optional[list[PropertyTuple]] = None
     metadata: dict[str, Any] = Field(default_factory=dict)
-    action: AnyAction = Field(default_factory=AcquireImage)
+    action: AnyAction = Field(default_factory=AcquireImage, discriminator="type")
     keep_shutter_open: bool = False
     reset_event_timer: bool = False
 

tests/test_sequence.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 
 from useq import (
     Channel,
@@ -23,6 +23,7 @@
     ZRangeAround,
     ZRelativePositions,
 )
+from useq._actions import CustomAction, HardwareAutofocus
 from useq._mda_event import SLMImage
 from useq._position import RelativePosition
 
@@ -354,7 +355,34 @@ def test_skip_channel_do_stack_no_zplan() -> None:
 
 def test_event_action_union() -> None:
     # test that action unions work
-    MDAEvent(action={"autofocus_device_name": "Z", "autofocus_motor_offset": 25})
+    event = MDAEvent(
+        action={
+            "type": "hardware_autofocus",
+            "autofocus_device_name": "Z",
+            "autofocus_motor_offset": 25,
+        }
+    )
+    assert isinstance(event.action, HardwareAutofocus)
+
+
+def test_custom_action() -> None:
+    event = MDAEvent(action={"type": "custom"})
+    assert isinstance(event.action, CustomAction)
+
+    event2 = MDAEvent(
+        action=CustomAction(
+            data={
+                "foo": "bar",
+                "alist": [1, 2, 3],
+                "nested": {"a": 1, "b": 2},
+                "nested_list": [{"a": 1}, {"b": 2}],
+            }
+        )
+    )
+    assert isinstance(event2.action, CustomAction)
+
+    with pytest.raises(ValidationError, match="must be JSON serializable"):
+        CustomAction(data={"not-serializable": lambda x: x})
 
 
 def test_keep_shutter_open() -> None: