Add GenericEncoder and Pydantic support (#63)

* Add GenericEncoder and Pydantic support

* Change default encoder clients

* Update type_util.py

* Add tests for pydantic models & increase coverage

* Fix mypy issues

* Add codegen for pydantic models

* Add is_allowed_type to encoder

* Add pydantic model to test_codegen

* Create test_pydantic_v1.py

* Add GenericEncoder test with pydantic v1

* Update Makefile

Author: martincpt

Makefile

@@ -29,7 +29,7 @@ format:
 install-lint:
 	python -m pip install --upgrade pip
 	pip install -r requirements.txt  # needed for pytype
-	pip install black isort flake8 pylint pytype mypy
+	pip install black isort flake8 pylint pytype mypy pydantic>=2.0
 
 lint:
 	flake8 ./zero

README.md

@@ -31,6 +31,7 @@
 *   Zero uses messages for communication and traditional **client-server** or **request-reply** pattern is supported.
 *   Support for both **async** and **sync**.
 *   The base server (ZeroServer) **utilizes all cpu cores**.
+*   Built-in support for Pydantic.
 *   **Code generation**! See [example](https://github.com/Ananto30/zero#code-generation-) 👇
 
 **Philosophy** behind Zero:

setup.py

@@ -22,4 +22,7 @@
     python_requires=">=3.8",
     package_dir={"": "."},
     install_requires=["pyzmq", "msgspec"],
+    extras_require={
+        "pydantic": ["pydantic"],  # Optional dependency
+    },
 )

tests/functional/codegen/test_codegen.py

@@ -11,6 +11,7 @@
 
 import msgspec
 from msgspec import Struct
+from pydantic import BaseModel
 
 from zero.codegen.codegen import CodeGen
 
@@ -63,6 +64,11 @@ class SimpleIntEnum(enum.IntEnum):
     TWO = 2
 
 
+class SimplePydanticModel(BaseModel):
+    a: int
+    b: str
+
+
 def func_none(arg: None) -> str:
     return "Received None"
 
@@ -213,6 +219,14 @@ def func_take_optional_child_dataclass_return_optional_child_complex_struct(
     return None
 
 
+def func_pydantic_model(arg: SimplePydanticModel) -> str:
+    return f"Received Pydantic model: {arg}"
+
+
+def func_return_pydantic_model() -> SimplePydanticModel:
+    return SimplePydanticModel(a=1, b="hello")
+
+
 class TestCodegen(unittest.TestCase):
     def setUp(self) -> None:
         self.maxDiff = None
@@ -250,6 +264,8 @@ def setUp(self) -> None:
             "func_msgspec_struct_complex": (func_msgspec_struct_complex, False),
             "func_child_complex_struct": (func_child_complex_struct, False),
             "func_return_complex_struct": (func_return_complex_struct, False),
+            "func_pydantic_model": (func_pydantic_model, False),
+            "func_return_pydantic_model": (func_return_pydantic_model, False),
         }
         self._rpc_input_type_map = {
             "func_none": None,
@@ -285,6 +301,8 @@ def setUp(self) -> None:
             "func_msgspec_struct_complex": ComplexStruct,
             "func_child_complex_struct": ChildComplexStruct,
             "func_return_complex_struct": None,
+            "func_pydantic_model": SimplePydanticModel,
+            "func_return_pydantic_model": None,
         }
         self._rpc_return_type_map = {
             "func_none": str,
@@ -320,6 +338,8 @@ def setUp(self) -> None:
             "func_msgspec_struct_complex": str,
             "func_child_complex_struct": str,
             "func_return_complex_struct": ComplexStruct,
+            "func_pydantic_model": str,
+            "func_return_pydantic_model": SimplePydanticModel,
         }
 
     def test_codegen(self):
@@ -335,6 +355,7 @@ def test_codegen(self):
 import enum
 import msgspec
 from msgspec import Struct
+from pydantic import BaseModel
 from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union
 import uuid
 
@@ -385,6 +406,11 @@ class ChildComplexStruct(ComplexStruct):
     i: str
 
 
+class SimplePydanticModel(BaseModel):
+    a: int
+    b: str
+
+
 
 class RpcClient:
     def __init__(self, zero_client: ZeroClient):
@@ -488,6 +514,12 @@ def func_child_complex_struct(self, arg: ChildComplexStruct) -> str:
 
     def func_return_complex_struct(self) -> ComplexStruct:
         return self._zero_client.call("func_return_complex_struct", None)
+
+    def func_pydantic_model(self, arg: SimplePydanticModel) -> str:
+        return self._zero_client.call("func_pydantic_model", arg)
+
+    def func_return_pydantic_model(self) -> SimplePydanticModel:
+        return self._zero_client.call("func_return_pydantic_model", None)
 """
         self.assertEqual(code, expected_code)
 

tests/functional/single_server/client_generation_test.py

@@ -23,6 +23,7 @@ def test_codegeneration():
 import decimal
 import enum
 import msgspec
+from pydantic import BaseModel
 from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union
 import uuid
 
@@ -49,6 +50,11 @@ class Dataclass:
     age: int
 
 
+class PydanticModel(BaseModel):
+    name: str
+    age: int
+
+
 class Message(msgspec.Struct):
     msg: str
     start_time: datetime
@@ -116,6 +122,9 @@ def echo_enum_int(self, msg: ColorInt) -> ColorInt:
     def echo_dataclass(self, msg: Dataclass) -> Dataclass:
         return self._zero_client.call("echo_dataclass", msg)
 
+    def echo_pydantic(self, msg: PydanticModel) -> PydanticModel:
+        return self._zero_client.call("echo_pydantic", msg)
+
     def echo_typing_tuple(self, msg: Tuple[int, str]) -> Tuple[int, str]:
         return self._zero_client.call("echo_typing_tuple", msg)
 

tests/functional/single_server/client_server_test.py

@@ -131,6 +131,13 @@ def test_echo_dataclass(zero_client):
     assert result == data
 
 
+# pydantic input
+def test_echo_pydantic(zero_client):
+    data = server.PydanticModel(name="John", age=30)
+    result = zero_client.call("echo_pydantic", data, return_type=server.PydanticModel)
+    assert result == data
+
+
 # typing.Tuple input
 def test_echo_typing_tuple(zero_client):
     assert zero_client.call(

tests/functional/single_server/server.py

@@ -9,6 +9,7 @@
 
 import jwt
 import msgspec
+from pydantic import BaseModel
 
 from zero import ZeroServer
 
@@ -155,6 +156,17 @@ def echo_dataclass(msg: Dataclass) -> Dataclass:
     return msg
 
 
+# pydantic input
+class PydanticModel(BaseModel):
+    name: str
+    age: int
+
+
+@app.register_rpc
+def echo_pydantic(msg: PydanticModel) -> PydanticModel:
+    return msg
+
+
 # typing.Tuple input
 @app.register_rpc
 def echo_typing_tuple(msg: typing.Tuple[int, str]) -> typing.Tuple[int, str]:

tests/functional/test_pydantic_v1.py

@@ -0,0 +1,41 @@
+import sys
+import importlib
+from typing import Iterator
+import pytest
+
+
+@pytest.fixture
+def patch_pydantic_to_v1(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    import pydantic.v1
+
+    # Patch sys.modules so any `import pydantic` gives you `pydantic.v1`
+    monkeypatch.setitem(sys.modules, "pydantic", pydantic.v1)
+    importlib.invalidate_caches()
+
+    yield
+
+    # Clean up after test
+    importlib.invalidate_caches()
+
+
+def test_module_with_pydantic_v1(patch_pydantic_to_v1: None) -> None:
+    # Re-import your module so it sees `pydantic` as v1
+    from zero.encoder import generic
+
+    importlib.reload(generic)
+
+    # Now run assertions that rely on v1 behavior
+    assert not generic.IS_PYDANTIC_V2
+
+    from pydantic import BaseModel
+
+    class TestModel(BaseModel):
+        name: str
+        age: int
+
+    encoder = generic.GenericEncoder()
+    model_instance = TestModel(name="Alice", age=30)
+    encoded_data = encoder.encode(model_instance)
+    decoded_instance = encoder.decode_type(encoded_data, TestModel)
+    assert decoded_instance.name == "Alice"
+    assert decoded_instance.age == 30

tests/requirements.txt

@@ -1,9 +1,10 @@
 pyzmq
 msgspec
+pydantic>=2.0
 pytest
 pytest-cov
 PyJWT
 pytest-asyncio
 tornado>=6.1
 requests
-pytest-timeout
+pytest-timeout

tests/unit/test_server.py

@@ -1,6 +1,6 @@
 import sys
 import unittest
-from typing import Any, Tuple
+from typing import Any, Tuple, Type
 from unittest.mock import patch
 
 # import pytest
@@ -71,9 +71,12 @@ def encode(self, message: Any) -> bytes:
             def decode(self, message: bytes) -> Any:
                 return message
 
-            def decode_type(self, message: bytes, typ: Any) -> Any:
+            def decode_type(self, message: bytes, typ: Type[Any]) -> Any:
                 return message
 
+            def is_allowed_type(self, typ: Type) -> bool:
+                return True
+
         encoder = CustomEncoder()
 
         server = ZeroServer(encoder=encoder)
@@ -94,9 +97,12 @@ def encode(self, message: Any) -> bytes:
             def decode(self, message: bytes) -> Any:
                 return message
 
-            def decode_type(self, message: bytes, typ: Any) -> Any:
+            def decode_type(self, message: bytes, typ: Type[Any]) -> Any:
                 return message
 
+            def is_allowed_type(self, typ: Type) -> bool:
+                return True
+
         encoder = CustomEncoder()
         port = 5562
 
@@ -118,9 +124,12 @@ def encode(self, message: Any) -> bytes:
             def decode(self, message: bytes) -> Any:
                 return message
 
-            def decode_type(self, message: bytes, typ: Any) -> Any:
+            def decode_type(self, message: bytes, typ: Type[Any]) -> Any:
                 return message
 
+            def is_allowed_type(self, typ: Type) -> bool:
+                return True
+
         encoder = CustomEncoder()
         host = "123.0.0.123"
 
@@ -142,9 +151,12 @@ def encode(self, message: Any) -> bytes:
             def decode(self, message: bytes) -> Any:
                 return message
 
-            def decode_type(self, message: bytes, typ: Any) -> Any:
+            def decode_type(self, message: bytes, typ: Type[Any]) -> Any:
                 return message
 
+            def is_allowed_type(self, typ: Type) -> bool:
+                return True
+
         encoder = CustomEncoder()
         host = "123.0.0.123"
         port = 5563
@@ -236,7 +248,7 @@ def test_register_rpc_with_long_name(self):
 
             @server.register_rpc
             def add_this_is_a_very_long_name_for_a_function_more_than_120_characters_ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff(
-                msg: Tuple[int, int]
+                msg: Tuple[int, int],
             ) -> int:
                 return msg[0] + msg[1]