Implement type converter and __get_pydantic_core_schema__

Author: dandavison

temporalio/contrib/pydantic/converter.py

@@ -1,31 +1,45 @@
+import inspect
 import json
-from typing import Any, Optional
+from typing import (
+    Any,
+    Optional,
+    Type,
+)
 
+import pydantic
 from pydantic.json import pydantic_encoder
+
+import temporalio.workflow
 from temporalio.api.common.v1 import Payload
 from temporalio.converter import (
     CompositePayloadConverter,
     DataConverter,
     DefaultPayloadConverter,
     JSONPlainPayloadConverter,
+    JSONTypeConverter,
 )
+from temporalio.worker.workflow_sandbox._restrictions import RestrictionContext
 
 
 class PydanticJSONPayloadConverter(JSONPlainPayloadConverter):
     """Pydantic JSON payload converter.
 
-    This extends the :py:class:`JSONPlainPayloadConverter` to override
-    :py:meth:`to_payload` using the Pydantic encoder.
+    Extends :py:class:`JSONPlainPayloadConverter` to override :py:meth:`to_payload` using
+    the Pydantic encoder. :py:meth:`from_payload` uses the parent implementation, with a
+    custom type converter.
     """
 
+    def __init__(self) -> None:
+        super().__init__(custom_type_converters=[PydanticModelTypeConverter()])
+
     def to_payload(self, value: Any) -> Optional[Payload]:
         """Convert all values with Pydantic encoder or fail.
 
         Like the base class, we fail if we cannot convert. This payload
         converter is expected to be the last in the chain, so it can fail if
         unable to convert.
         """
-        # We let JSON conversion errors be thrown to caller
+        # Let JSON conversion errors be thrown to caller
         return Payload(
             metadata={"encoding": self.encoding.encode()},
             data=json.dumps(
@@ -34,17 +48,47 @@ def to_payload(self, value: Any) -> Optional[Payload]:
         )
 
 
+class PydanticModelTypeConverter(JSONTypeConverter):
+    def to_typed_value(self, hint: Type, value: Any) -> Any:
+        if not inspect.isclass(hint) or not issubclass(hint, pydantic.BaseModel):
+            return JSONTypeConverter.Unhandled
+        model = hint
+        if not isinstance(value, dict):
+            raise TypeError(
+                f"Cannot convert to {model}, value is {type(value)} not dict"
+            )
+        if temporalio.workflow.unsafe.in_sandbox():
+            # Unwrap proxied model field types so that Pydantic can call their constructors
+            model = pydantic.create_model(
+                model.__name__,
+                **{  # type: ignore
+                    name: (RestrictionContext.unwrap_if_proxied(f.annotation), f)
+                    for name, f in model.model_fields.items()
+                },
+            )
+        if hasattr(model, "model_validate"):
+            return model.model_validate(value)
+        elif hasattr(model, "parse_obj"):
+            # Pydantic v1
+            return model.parse_obj(value)
+        else:
+            raise ValueError(
+                f"{model} is a Pydantic model but does not have a `model_validate` or `parse_obj` method"
+            )
+
+
 class PydanticPayloadConverter(CompositePayloadConverter):
     """Payload converter that replaces Temporal JSON conversion with Pydantic
     JSON conversion.
     """
 
     def __init__(self) -> None:
+        json_payload_converter = PydanticJSONPayloadConverter()
         super().__init__(
             *(
                 c
                 if not isinstance(c, JSONPlainPayloadConverter)
-                else PydanticJSONPayloadConverter()
+                else json_payload_converter
                 for c in DefaultPayloadConverter.default_encoding_payload_converters
             )
         )

temporalio/converter.py

@@ -558,9 +558,10 @@ def encoding(self) -> str:
     def to_payload(self, value: Any) -> Optional[temporalio.api.common.v1.Payload]:
         """See base class."""
         # Check for pydantic then send warning
+        # TODO (dan): update
         if hasattr(value, "parse_obj"):
             warnings.warn(
-                "If you're using pydantic model, refer to https://github.com/temporalio/samples-python/tree/main/pydantic_converter for better support"
+                "If you're using a pydantic model, refer to https://github.com/temporalio/samples-python/tree/main/pydantic_converter for better support"
             )
         # We let JSON conversion errors be thrown to caller
         return temporalio.api.common.v1.Payload(
@@ -1522,18 +1523,6 @@ def value_to_type(
         # TODO(cretz): Want way to convert snake case to camel case?
         return hint(**field_values)
 
-    # If there is a @staticmethod or @classmethod parse_obj, we will use it.
-    # This covers Pydantic models.
-    parse_obj_attr = inspect.getattr_static(hint, "parse_obj", None)
-    if isinstance(parse_obj_attr, classmethod) or isinstance(
-        parse_obj_attr, staticmethod
-    ):
-        if not isinstance(value, dict):
-            raise TypeError(
-                f"Cannot convert to {hint}, value is {type(value)} not dict"
-            )
-        return getattr(hint, "parse_obj")(value)
-
     # IntEnum
     if inspect.isclass(hint) and issubclass(hint, IntEnum):
         if not isinstance(value, int):

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__