Implement __get_pydantic_core_schema__

Author: dandavison

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(object.__getattribute__(source_type, "__temporal_state").obj)
+        )
+
     def __init__(self, *args, **kwargs) -> None:
         # When we instantiate this class, we have the signature of:
         #   __init__(
@@ -973,6 +984,8 @@ def __init__(self, *args, **kwargs) -> None:
     def __getattribute__(self, __name: str) -> Any:
         state = _RestrictionState.from_proxy(self)
         _trace("__getattribute__ %s on %s", __name, state.name)
+        if __name == "__get_pydantic_core_schema__":
+            return object.__getattribute__(self, "__get_pydantic_core_schema__")
         # We do not restrict __spec__ or __name__
         if __name != "__spec__" and __name != "__name__":
             state.assert_child_not_restricted(__name)