Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions src/python_workflow_definition/models.py
Original file line number Diff line number Diff line change
@@ -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,
TypeVar,
)
from typing_extensions import TypeAliasType
from pydantic import BaseModel, Field, field_validator, field_serializer
from pydantic import ValidationError
import json
Expand All @@ -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."""
Expand All @@ -33,7 +48,7 @@ class PythonWorkflowDefinitionInputNode(PythonWorkflowDefinitionBaseNode):

type: Literal["input"]
name: str
value: Optional[Any] = None
value: Optional[AllowableDefaults] = None


class PythonWorkflowDefinitionOutputNode(PythonWorkflowDefinitionBaseNode):
Expand Down
50 changes: 50 additions & 0 deletions tests/test_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,11 @@
INTERNAL_DEFAULT_HANDLE,
)


class _NoTrivialSerialization:
pass


class TestModels(unittest.TestCase):
def setUp(self):
self.valid_workflow_dict = {
Expand Down Expand Up @@ -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)
Expand Down
Loading