feat: add pydantic codec testing

Author: brian14708

pyproject.toml

@@ -14,6 +14,7 @@ build-backend = "uv_build"
 
 [dependency-groups]
 dev = [
+    "pydantic>=2.11.9",
     "pytest>=7.0",
     "pytest-asyncio>=0.21",
 ]
@@ -46,8 +47,12 @@ preview = true
 future-annotations = true
 
 [tool.ruff.lint.flake8-type-checking]
+runtime-evaluated-base-classes = ["typing_extensions.TypedDict"]
 runtime-evaluated-decorators = ["duron.fn"]
 
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+"typing.TypedDict".msg = "Use typing_extensions.TypedDict instead."
+
 [tool.ruff.format]
 preview = true
 

src/duron/codec.py

@@ -1,17 +1,36 @@
 from __future__ import annotations
 
 import inspect
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, final
+from typing import TYPE_CHECKING, TypeGuard, cast, final
 
-from typing_extensions import Protocol, override
-
-from duron.log import is_json_value
+from typing_extensions import TypeAliasType, override
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from duron.log import JSONValue
+    JSONValue = (
+        dict[str, "JSONValue"] | list["JSONValue"] | str | int | float | bool | None
+    )
+else:
+    JSONValue = TypeAliasType(
+        "JSONValue",
+        "dict[str, JSONValue] | list[JSONValue] | str | int | float | bool | None",
+    )
+
+
+def is_json_value(x: object) -> TypeGuard[JSONValue]:
+    if x is None or isinstance(x, (bool, int, float, str)):
+        return True
+    if isinstance(x, list):
+        return all(is_json_value(item) for item in cast("list[object]", x))
+    if isinstance(x, dict):
+        return all(
+            isinstance(k, str) and is_json_value(v)
+            for k, v in cast("dict[object, object]", x).items()
+        )
+    return False
 
 
 @dataclass(slots=True)
@@ -21,31 +40,15 @@ class FunctionType:
     parameter_types: dict[str, type | None]
 
 
-class Codec(Protocol):
+class Codec(ABC):
+    @abstractmethod
     def encode_json(self, result: object, /) -> JSONValue: ...
+
+    @abstractmethod
     def decode_json(
         self, encoded: JSONValue, expected_type: type | None, /
     ) -> object: ...
 
-    def inspect_function(
-        self,
-        fn: Callable[..., object],
-    ) -> FunctionType: ...
-
-
-@final
-class DefaultCodec(Codec):
-    @override
-    def encode_json(self, result: object) -> JSONValue:
-        if is_json_value(result):
-            return result
-        raise TypeError(f"Result is not JSON-serializable: {result!r}")
-
-    @override
-    def decode_json(self, encoded: JSONValue, _expected_type: type | None) -> object:
-        return encoded
-
-    @override
     def inspect_function(
         self,
         fn: Callable[..., object],
@@ -76,3 +79,16 @@ def inspect_function(
             parameters=parameter_names,
             parameter_types=parameter_types,
         )
+
+
+@final
+class DefaultCodec(Codec):
+    @override
+    def encode_json(self, result: object) -> JSONValue:
+        if is_json_value(result):
+            return result
+        raise TypeError(f"Result is not JSON-serializable: {result!r}")
+
+    @override
+    def decode_json(self, encoded: JSONValue, _expected_type: type | None) -> object:
+        return encoded

src/duron/contrib/codecs.py

@@ -4,23 +4,21 @@
 import pickle
 from typing import TYPE_CHECKING
 
-from duron.codec import DefaultCodec
+from typing_extensions import override
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
+from duron.codec import Codec
 
-    from duron.codec import FunctionType
-    from duron.log import JSONValue
+if TYPE_CHECKING:
+    from duron.codec import JSONValue
 
 
-class PickleCodec:
+class PickleCodec(Codec):
+    @override
     def encode_json(self, result: object) -> str:
         return base64.b64encode(pickle.dumps(result)).decode()
 
+    @override
     def decode_json(self, encoded: JSONValue, _expected_type: type | None) -> object:
         if not isinstance(encoded, str):
             raise TypeError(f"Expected a string, got {type(encoded).__name__}")
         return pickle.loads(base64.b64decode(encoded.encode()))
-
-    def inspect_function(self, fn: Callable[..., object]) -> FunctionType:
-        return DefaultCodec().inspect_function(fn)

src/duron/log.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Generic, Protocol, cast
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, Literal
 
-from typing_extensions import TypedDict, TypeVar
+from typing_extensions import NotRequired, TypedDict, TypeVar
+
+from duron.codec import JSONValue
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
-    from typing import Literal, TypeGuard
-
-    from typing_extensions import NotRequired
+    from typing import TypeGuard
 
 
 _TOffset = TypeVar("_TOffset")
@@ -78,31 +79,20 @@ def is_entry(entry: Entry | AnyEntry) -> TypeGuard[Entry]:
     }
 
 
-JSONValue = dict[str, "JSONValue"] | list["JSONValue"] | str | int | float | bool | None
-
-
-def is_json_value(x: object) -> TypeGuard[JSONValue]:
-    if x is None or isinstance(x, (bool, int, float, str)):
-        return True
-    if isinstance(x, list):
-        return all(is_json_value(item) for item in cast("list[object]", x))
-    if isinstance(x, dict):
-        return all(
-            isinstance(k, str) and is_json_value(v)
-            for k, v in cast("dict[object, object]", x).items()
-        )
-    return False
-
-
-class LogStorage(Protocol, Generic[_TOffset, _TLease]):
+class LogStorage(ABC, Generic[_TOffset, _TLease]):
+    @abstractmethod
     def stream(
         self, start: _TOffset | None, live: bool, /
     ) -> AsyncGenerator[tuple[_TOffset, AnyEntry], None]: ...
 
+    @abstractmethod
     async def acquire_lease(self) -> _TLease: ...
 
+    @abstractmethod
     async def release_lease(self, token: _TLease, /): ...
 
+    @abstractmethod
     async def append(self, token: _TLease, entry: Entry, /): ...
 
+    @abstractmethod
     async def flush(self, token: _TLease, /): ...

src/duron/task.py

@@ -11,13 +11,14 @@
     Any,
     Generic,
     ParamSpec,
-    TypedDict,
     TypeVar,
     cast,
     final,
 )
 
-from duron.codec import Codec
+from typing_extensions import TypedDict
+
+from duron.codec import Codec, JSONValue
 from duron.context import Context
 from duron.event_loop import EventLoop, create_loop
 from duron.log import is_entry
@@ -33,7 +34,6 @@
     from duron.log import (
         Entry,
         ErrorInfo,
-        JSONValue,
         LogStorage,
         PromiseCompleteEntry,
     )
@@ -73,14 +73,16 @@ def __init__(
         self._run: _TaskRun | None = None
 
     async def start(self, *args: _P.args, **kwargs: _P.kwargs) -> None:
+        def get_init() -> TaskInitParams:
+            return {
+                "version": _CURRENT_VERSION,
+                "args": [codec.encode_json(arg) for arg in args],
+                "kwargs": {k: codec.encode_json(v) for k, v in kwargs.items()},
+            }
+
         codec = self._task_fn.codec
-        init: TaskInitParams = {
-            "version": _CURRENT_VERSION,
-            "args": [codec.encode_json(arg) for arg in args],
-            "kwargs": {k: codec.encode_json(v) for k, v in kwargs.items()},
-        }
         type_info = codec.inspect_function(self._task_fn.fn)
-        task_prelude = _task_prelude(self._task_fn, type_info, lambda: init)
+        task_prelude = _task_prelude(self._task_fn, type_info, get_init)
         self._run = _TaskRun(
             TaskRun(task=task_prelude, return_type=type_info.return_type),
             self._log,
@@ -162,7 +164,7 @@ def __init__(
         self._pending_ops: set[bytes] = set()
         self._now = 0
         self._offset: object | None = None
-        self._tasks: dict[str, tuple[asyncio.Task[None], type | None]] = {}
+        self._tasks: dict[str, tuple[asyncio.Future[None], type | None]] = {}
 
     def now(self) -> int:
         if self._running:
@@ -256,25 +258,34 @@ async def handle_message(self, e: Entry) -> None:
             pending_info = self._pending_task.pop(e["promise_id"], None)
             task_info = self._tasks.pop(e["promise_id"], None)
 
-            # Get return type from either pending or running task
+            id = _decode_id(e["promise_id"])
+            if id not in self._pending_ops:
+                return
+
             return_type = None
             if pending_info is not None:
                 _, return_type = pending_info
             elif task_info is not None:
                 _, return_type = task_info
+            else:
+                print(e)
+                raise AssertionError("unreachable")
 
-            id = _decode_id(e["promise_id"])
             if "error" in e:
                 self._loop.post_completion(
                     id,
                     exception=_decode_error(e["error"]),
                 )
                 self._pending_ops.discard(id)
             elif "result" in e:
-                self._loop.post_completion(
-                    id,
-                    result=self._codec.decode_json(e["result"], return_type),
-                )
+                try:
+                    result = self._codec.decode_json(e["result"], return_type)
+                    self._loop.post_completion(id, result=result)
+                except BaseException as exc:
+                    self._loop.post_completion(
+                        id,
+                        exception=exc,
+                    )
                 self._pending_ops.discard(id)
             else:
                 raise ValueError(f"Invalid promise/complete entry: {e!r}")
@@ -339,6 +350,8 @@ def done(f: OpFuture) -> None:
                     "type": "promise/create",
                 })
 
+                fut_host: asyncio.Future[None] = asyncio.Future()
+
                 async def cb() -> None:
                     entry: PromiseCompleteEntry = {
                         "ts": _encode_timestamp(self.now()),
@@ -353,9 +366,15 @@ async def cb() -> None:
                         entry["error"] = _encode_error(e)
                     finally:
                         await self.enqueue_log(entry, True)
+                        fut_host.set_result(None)
 
                 _ = self._loop.create_task(cb())
 
+                self._tasks[_encode_id(id)] = (
+                    fut_host,
+                    op.return_type,
+                )
+
             case _:
                 raise NotImplementedError(f"Unsupported op: {op!r}")
 

tests/test_pydantic.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import pytest
+from pydantic import BaseModel, TypeAdapter
+from typing_extensions import override
+
+from duron import fn
+from duron.codec import Codec
+from duron.context import Context
+from duron.contrib.storage import MemoryLogStorage
+
+if TYPE_CHECKING:
+    from duron.codec import JSONValue
+
+
+class PydanticPoint(BaseModel):
+    x: int
+    y: int
+
+
+class PydanticCodec(Codec):
+    @override
+    def encode_json(self, result: object) -> JSONValue:
+        return cast(
+            "JSONValue", TypeAdapter(type(result)).dump_python(result, mode="json")
+        )
+
+    @override
+    def decode_json(self, encoded: JSONValue, expected_type: type | None) -> object:
+        return cast("object", TypeAdapter(expected_type).validate_python(encoded))
+
+
+@pytest.mark.asyncio
+async def test_pydantic_serialize():
+    @fn(codec=PydanticCodec())
+    async def activity(ctx: Context) -> PydanticPoint:
+        def new_pt() -> PydanticPoint:
+            return PydanticPoint(x=1, y=2)
+
+        pt = await ctx.run(new_pt)
+        return PydanticPoint(x=pt.x + 5, y=pt.y + 10)
+
+    log = MemoryLogStorage()
+    async with activity.create_task(log) as t:
+        await t.start()
+        a = await t.wait()
+        assert type(a) is PydanticPoint
+        assert a.x == 6 and a.y == 12