instancify_args: support Optional dataclasses

src/saturn_engine/core/pipeline.py

@@ -1,3 +1,4 @@
+import typing
 from typing import Any
 from typing import Callable
 from typing import cast
@@ -48,8 +49,21 @@ def instancify_args(args: dict[str, object], pipeline: Callable) -> None:
             if not isinstance(data, dict):
                 continue
 
-            if dataclasses.is_dataclass(parameter.annotation):
-                schema = schema_for(parameter.annotation)
+            target_type = parameter.annotation
+
+            # Support Optional[dataclass], which is very common.
+            # This does not support Union of two dataclasses yet.
+            if typing.get_origin(target_type) is typing.Union:
+                target_args = [
+                    arg
+                    for arg in typing.get_args(target_type)
+                    if not isinstance(None, arg)
+                ]
+                if len(target_args) == 1:
+                    target_type = target_args[0]
+
+            if dataclasses.is_dataclass(target_type):
+                schema = schema_for(target_type)
                 args[parameter.name] = schema.load(data)
 
 

tests/core/test_pipeline.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+import dataclasses
+
+from saturn_engine.core.pipeline import PipelineInfo
+
+
+def test_instancify_args_optional() -> None:
+    @dataclasses.dataclass
+    class A:
+        field: str
+
+    def pipeline(a: Optional[A]) -> None:
+        pass
+
+    args: dict[str, object] = {"a": {"field": "value"}}
+
+    PipelineInfo.instancify_args(
+        args=args,
+        pipeline=pipeline,
+    )
+
+    assert args == {"a": A(field="value")}