WIP - support pydantic

README.md

@@ -262,6 +262,4 @@ ZnFlow includes tests to ensure compatibility with:
 - `dataclasses`
 - `ZnInit`
 - `attrs`
-
-It is currently **not** compatible with pydantic. I don't know what pydantic
-does internally and wasn't able to find a workaround.
+- `pydantic` (experimental)

pyproject.toml

@@ -22,6 +22,7 @@ pytest = "^7"
 coverage = "^6"
 zninit = "^0.1"
 attrs = "^22"
+pydantic = "^2.6.2"
 
 [tool.poetry.group.notebook.dependencies]
 jupyterlab = "^3"

tests/test_node.py

@@ -1,6 +1,7 @@
 import dataclasses
 
 import attrs
+import pydantic
 import pytest
 import zninit
 
@@ -58,6 +59,17 @@ def output(self):
         return znflow.get_attribute(self, "value")
 
 
+class PydanticNode(pydantic.BaseModel, znflow.Node):
+    value: pydantic.SkipValidation[int]
+
+    def run(self):
+        self.value += 1
+
+    @property
+    def output(self):
+        return znflow.get_attribute(self, "value")
+
+
 @znflow.nodify
 def add(value):
     return value
@@ -68,15 +80,20 @@ def compute_sum(*args):
     return sum(args)
 
 
-@pytest.mark.parametrize("cls", [PlainNode, DataclassNode, ZnInitNode, add, AttrsNode])
+@pytest.mark.parametrize(
+    "cls", [PlainNode, DataclassNode, ZnInitNode, add, AttrsNode, PydanticNode]
+)
 def test_Node_init(cls):
-    with pytest.raises((TypeError, AttributeError)):
+    with pytest.raises((TypeError, AttributeError, pydantic.ValidationError)):
         # TODO only raise TypeError and not AttributeError when TypeError is expected.
+        # Pydantic does not raise TypeError, but ValidationError
         with znflow.DiGraph():
             cls()
 
 
-@pytest.mark.parametrize("cls", [PlainNode, DataclassNode, ZnInitNode, add, AttrsNode])
+@pytest.mark.parametrize(
+    "cls", [PlainNode, DataclassNode, ZnInitNode, add, AttrsNode, PydanticNode]
+)
 def test_Node(cls):
     with znflow.DiGraph() as graph:
         node = cls(value=42)
@@ -94,8 +111,12 @@ def test_Node(cls):
     assert graph.nodes[node.uuid]["value"] is node
 
 
-@pytest.mark.parametrize("cls2", [PlainNode, DataclassNode, ZnInitNode, AttrsNode])
-@pytest.mark.parametrize("cls1", [PlainNode, DataclassNode, ZnInitNode, AttrsNode])
+@pytest.mark.parametrize(
+    "cls2", [PlainNode, DataclassNode, ZnInitNode, AttrsNode, PydanticNode]
+)
+@pytest.mark.parametrize(
+    "cls1", [PlainNode, DataclassNode, ZnInitNode, AttrsNode, PydanticNode]
+)
 def test_ConnectionNodeNode(cls1, cls2):
     with znflow.DiGraph() as graph:
         node1 = cls1(value=42)
@@ -132,7 +153,9 @@ def test_ConnectionNodifyNodify(cls1, cls2):
 
 
 @pytest.mark.parametrize("cls1", [add])
-@pytest.mark.parametrize("cls2", [PlainNode, DataclassNode, ZnInitNode, AttrsNode])
+@pytest.mark.parametrize(
+    "cls2", [PlainNode, DataclassNode, ZnInitNode, AttrsNode, PydanticNode]
+)
 def test_ConnectionNodeNodify(cls1, cls2):
     with znflow.DiGraph() as graph:
         node1 = cls1(value=42)
@@ -151,7 +174,9 @@ def test_ConnectionNodeNodify(cls1, cls2):
 
 
 @pytest.mark.parametrize("cls2", [add])
-@pytest.mark.parametrize("cls1", [PlainNode, DataclassNode, ZnInitNode, AttrsNode])
+@pytest.mark.parametrize(
+    "cls1", [PlainNode, DataclassNode, ZnInitNode, AttrsNode, PydanticNode]
+)
 def test_ConnectionNodifyNode(cls1, cls2):
     with znflow.DiGraph() as graph:
         node1 = cls1(value=42)
@@ -167,7 +192,9 @@ def test_ConnectionNodifyNode(cls1, cls2):
 
 
 @pytest.mark.parametrize("cls2", [compute_sum])
-@pytest.mark.parametrize("cls1", [PlainNode, DataclassNode, ZnInitNode, AttrsNode])
+@pytest.mark.parametrize(
+    "cls1", [PlainNode, DataclassNode, ZnInitNode, AttrsNode, PydanticNode]
+)
 def test_ConnectionNodifyMultiNode(cls1, cls2):
     with znflow.DiGraph() as graph:
         node1 = cls1(value=42)
@@ -190,7 +217,9 @@ def test_ConnectionNodifyMultiNode(cls1, cls2):
 
 
 @pytest.mark.parametrize("cls1", [compute_sum])
-@pytest.mark.parametrize("cls2", [PlainNode, DataclassNode, ZnInitNode, AttrsNode])
+@pytest.mark.parametrize(
+    "cls2", [PlainNode, DataclassNode, ZnInitNode, AttrsNode, PydanticNode]
+)
 def test_ConnectionNodeMultiNodify(cls1, cls2):
     with znflow.DiGraph() as graph:
         node1 = cls1(42)

znflow/base.py

@@ -113,7 +113,9 @@ class NodeBaseMixin:
         "_graph_",
         "uuid",
         "_uuid",
-    ]  # TODO consider adding regex patterns
+        "model_fields",  # pydantic
+        "model_computed_fields",  # pydantic
+    ]
 
     @property
     def uuid(self):

znflow/graph.py

@@ -108,9 +108,11 @@ def _update_node_attributes(self, node_instance: Node, updater) -> None:
                 value, node_instance=node_instance, attribute=attribute
             )
 
-    def add_node(self, node_for_adding, **attr):
+    def add_node(self, node_for_adding, this_uuid=None, **attr):
         if isinstance(node_for_adding, NodeBaseMixin):
-            super().add_node(node_for_adding.uuid, value=node_for_adding, **attr)
+            if this_uuid is None:
+                this_uuid = node_for_adding.uuid
+            super().add_node(this_uuid, value=node_for_adding, **attr)
         else:
             raise ValueError(f"Only Nodes are supported, found '{node_for_adding}'.")