From 34d1ae2534afb132c6f44d3a7ebcd58c0b0f97cf Mon Sep 17 00:00:00 2001 From: Dan Davison Date: Mon, 3 Feb 2025 23:48:11 -0500 Subject: [PATCH] Only use pydantic in sandbox if it can be imported --- .../worker/workflow_sandbox/_restrictions.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/temporalio/worker/workflow_sandbox/_restrictions.py b/temporalio/worker/workflow_sandbox/_restrictions.py index 321f310a..8e8b70cd 100644 --- a/temporalio/worker/workflow_sandbox/_restrictions.py +++ b/temporalio/worker/workflow_sandbox/_restrictions.py @@ -32,8 +32,11 @@ cast, ) -from pydantic import GetCoreSchemaHandler -from pydantic_core import CoreSchema, core_schema +try: + import pydantic + import pydantic_core +except ImportError: + pydantic = None # type: ignore import temporalio.workflow @@ -986,7 +989,7 @@ 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__": + if pydantic and __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) @@ -1037,14 +1040,17 @@ 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)) - ) + if pydantic: + # Instruct pydantic to use the proxied type when determining the schema + @classmethod + def __get_pydantic_core_schema__( + cls, + source_type: Any, + handler: pydantic.GetCoreSchemaHandler, # type: ignore + ) -> pydantic_core.CoreSchema: + return pydantic_core.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