Cleanup

Author: dandavison

temporalio/contrib/pydantic.py

@@ -23,7 +23,7 @@
     from pydantic_core import to_jsonable_python
 except ImportError:
     # pydantic v1
-    from pydantic.json import pydantic_encoder as to_jsonable_python
+    from pydantic.json import pydantic_encoder as to_jsonable_python  # type: ignore
 
 import temporalio.workflow
 from temporalio.converter import (
@@ -36,9 +36,16 @@
 )
 from temporalio.worker.workflow_sandbox._restrictions import RestrictionContext
 
+# Note that in addition to the implementation in this module, _RestrictedProxy
+# implements __get_pydantic_core_schema__ so that pydantic unwraps proxied types
+# when determining the schema.
+
 
 class PydanticModelTypeConverter(JSONTypeConverter):
+    """Type converter for pydantic model instances."""
+
     def to_typed_value(self, hint: Type, value: Any) -> Any:
+        """Convert dict value to pydantic model instance of the specified type"""
         if not inspect.isclass(hint) or not issubclass(hint, pydantic.BaseModel):
             return JSONTypeConverter.Unhandled
         model = hint
@@ -67,20 +74,29 @@ def to_typed_value(self, hint: Type, value: Any) -> Any:
 
 
 class PydanticJSONEncoder(AdvancedJSONEncoder):
+    """JSON encoder for python objects containing pydantic model instances."""
+
     def default(self, o: Any) -> Any:
+        """Convert object to jsonable python.
+
+        See :py:meth:`json.JSONEncoder.default`.
+        """
         if isinstance(o, pydantic.BaseModel):
             return to_jsonable_python(o)
         return super().default(o)
 
 
 class PydanticPayloadConverter(CompositePayloadConverter):
-    """Pydantic payload converter.
+    """Payload converter for payloads containing pydantic model instances.
 
-    Payload converter that replaces the default JSON conversion with Pydantic
-    JSON conversion.
+    JSON conversion is replaced with a converter that uses
+    :py:class:`PydanticJSONEncoder` to convert the python object to JSON, and
+    :py:class:`PydanticModelTypeConverter` to convert raw python values to
+    pydantic model instances.
     """
 
     def __init__(self) -> None:
+        """Initialize object"""
         json_payload_converter = JSONPlainPayloadConverter(
             encoder=PydanticJSONEncoder,
             custom_type_converters=[PydanticModelTypeConverter()],
@@ -98,7 +114,7 @@ def __init__(self) -> None:
 pydantic_data_converter = DataConverter(
     payload_converter_class=PydanticPayloadConverter
 )
-"""Data converter for Pydantic models.
+"""Data converter for payloads containing pydantic model instances.
 
-To use, pass this as the ``data_converter`` argument to :py:class:`temporalio.client.Client`
+To use, pass as the ``data_converter`` argument of :py:class:`temporalio.client.Client`
 """

temporalio/worker/workflow_sandbox/_restrictions.py

@@ -951,14 +951,6 @@ 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__(
@@ -1033,6 +1025,15 @@ def __getitem__(self, key: Any) -> Any:
                 )
         return ret
 
+    # Instruct pydantic to use the proxied type when determining the schema
+    @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))
+        )
+
     __doc__ = _RestrictedProxyLookup(  # type: ignore
         class_value=__doc__, fallback_func=lambda self: type(self).__doc__, is_attr=True
     )