Implement pydantic from_payload and __get_pydantic_core_schema__

Author: dandavison

temporalio/contrib/pydantic/converter.py

@@ -1,14 +1,21 @@
 import json
-from typing import Any, Optional
+import typing
+from typing import Any, Optional, Type
 
+from pydantic import BaseModel, create_model
 from pydantic.json import pydantic_encoder
+
+import temporalio.workflow
 from temporalio.api.common.v1 import Payload
 from temporalio.converter import (
     CompositePayloadConverter,
     DataConverter,
     DefaultPayloadConverter,
     JSONPlainPayloadConverter,
 )
+from temporalio.worker.workflow_sandbox._restrictions import (
+    RestrictionContext,
+)
 
 
 class PydanticJSONPayloadConverter(JSONPlainPayloadConverter):
@@ -33,6 +40,26 @@ def to_payload(self, value: Any) -> Optional[Payload]:
             ).encode(),
         )
 
+    def from_payload(self, payload: Payload, type_hint: Optional[Type] = None) -> Any:
+        data = json.loads(payload.data.decode())
+        if type_hint and typing.get_origin(type_hint) is list:
+            assert isinstance(data, list), "Expected list"
+            [type_hint] = typing.get_args(type_hint)
+            assert type_hint is not None, "Expected type hint"
+            assert issubclass(type_hint, BaseModel), "Expected BaseModel"
+            if temporalio.workflow.unsafe.in_sandbox():
+                type_hint = _unwrap_restricted_fields(type_hint)
+
+            return [self._from_dict(d, type_hint) for d in data]
+        return self._from_dict(data, type_hint)
+
+    def _from_dict(self, data: dict, type_hint: Optional[Type]) -> Any:
+        assert isinstance(data, dict), "Expected dict"
+        if type_hint and hasattr(type_hint, "validate"):
+            return type_hint.validate(data)
+
+        return data
+
 
 class PydanticPayloadConverter(CompositePayloadConverter):
     """Payload converter that replaces Temporal JSON conversion with Pydantic
@@ -54,3 +81,13 @@ def __init__(self) -> None:
     payload_converter_class=PydanticPayloadConverter
 )
 """Data converter using Pydantic JSON conversion."""
+
+
+def _unwrap_restricted_fields(
+    model: Type[BaseModel],
+) -> Type[BaseModel]:
+    fields = {
+        name: (RestrictionContext.unwrap_if_proxied(f.annotation), f)
+        for name, f in model.model_fields.items()
+    }
+    return create_model(model.__name__, **fields)  # type: ignore

temporalio/worker/workflow_sandbox/_restrictions.py

@@ -31,6 +31,9 @@
     cast,
 )
 
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import CoreSchema, core_schema
+
 import temporalio.workflow
 
 logger = logging.getLogger(__name__)
@@ -948,6 +951,14 @@ def _is_restrictable(v: Any) -> bool:
 
 
 class _RestrictedProxy:
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        return core_schema.no_info_after_validator_function(
+            cls, handler(RestrictionContext.unwrap_if_proxied(source_type))
+        )
+
     def __init__(self, *args, **kwargs) -> None:
         # When we instantiate this class, we have the signature of:
         #   __init__(
@@ -971,6 +982,8 @@ def __init__(self, *args, **kwargs) -> None:
             _trace("__init__ unrecognized with args %s", args)
 
     def __getattribute__(self, __name: str) -> Any:
+        if __name == "__get_pydantic_core_schema__":
+            return object.__getattribute__(self, "__get_pydantic_core_schema__")
         state = _RestrictionState.from_proxy(self)
         _trace("__getattribute__ %s on %s", __name, state.name)
         # We do not restrict __spec__ or __name__