Skip to content

Commit

Permalink
Only use pydantic in sandbox if it can be imported
Browse files Browse the repository at this point in the history
  • Loading branch information
dandavison committed Feb 4, 2025
1 parent d3d9ed3 commit 34d1ae2
Showing 1 changed file with 17 additions and 11 deletions.
28 changes: 17 additions & 11 deletions temporalio/worker/workflow_sandbox/_restrictions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 34d1ae2

Please sign in to comment.