Allow immutable json primitives recursively

Signed-off-by: liamhuber <liamhuber@greyhavensolutions.com>

Author: liamhuber

src/python_workflow_definition/models.py

@@ -1,5 +1,15 @@
 from pathlib import Path
-from typing import List, Union, Optional, Literal, Any, Annotated, Type, TypeVar
+from typing import (
+    List,
+    Union,
+    Optional,
+    Literal,
+    Any,
+    Annotated,
+    Type,
+    TypeAliasType,
+    TypeVar,
+)
 from pydantic import BaseModel, Field, field_validator, field_serializer
 from pydantic import ValidationError
 import json
@@ -18,6 +28,11 @@
     "PythonWorkflowDefinitionWorkflow",
 )
 
+JsonPrimitive = Union[str, int, float, bool, None]
+AllowableDefaults = TypeAliasType(
+    "AllowableDefaults", "Union[JsonPrimitive, tuple[AllowableDefaults, ...]]"
+)
+
 
 class PythonWorkflowDefinitionBaseNode(BaseModel):
     """Base model for all node types, containing common fields."""
@@ -33,7 +48,7 @@ class PythonWorkflowDefinitionInputNode(PythonWorkflowDefinitionBaseNode):
 
     type: Literal["input"]
     name: str
-    value: Optional[Any] = None
+    value: Optional[AllowableDefaults] = None
 
 
 class PythonWorkflowDefinitionOutputNode(PythonWorkflowDefinitionBaseNode):

tests/test_models.py

@@ -12,6 +12,11 @@
     INTERNAL_DEFAULT_HANDLE,
 )
 
+
+class _NoTrivialSerialization:
+    pass
+
+
 class TestModels(unittest.TestCase):
     def setUp(self):
         self.valid_workflow_dict = {
@@ -40,6 +45,51 @@ def test_input_node(self):
         )
         self.assertEqual(node_with_value.value, 42)
 
+    def test_input_node_valid_values(self):
+        good_values = (
+            1,
+            1.1,
+            "string",
+            True,
+            None,
+            (1, 2),
+            (("recursive", "tuple"), (True, False)),
+        )
+        for value in good_values:
+            with self.subTest(value=value):
+                model = PythonWorkflowDefinitionInputNode.model_validate(
+                    {
+                        "id": 0,
+                        "type": "input",
+                        "name": "x",
+                        "value": value,
+                    }
+                )
+                self.assertEqual(
+                    value,
+                    PythonWorkflowDefinitionInputNode.model_validate(
+                        model.model_dump(mode="json")
+                    ).value
+                )
+
+
+    def test_input_node_invalid_value_raises(self):
+        bad_values = (
+            {"mutable": "thing"},
+            _NoTrivialSerialization(),
+        )
+        for value in bad_values:
+            with self.subTest(value=value):
+                with self.assertRaises(ValidationError):
+                    PythonWorkflowDefinitionInputNode.model_validate(
+                        {
+                            "id": 0,
+                            "type": "input",
+                            "name": "x",
+                            "value": value,
+                        }
+                    )
+
     def test_output_node(self):
         node = PythonWorkflowDefinitionOutputNode(id=1, type="output", name="test_output")
         self.assertEqual(node.id, 1)