1from __future__ import annotations 
    2 
    3from collections.abc import Iterator, Mapping 
    4from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, TypeVar, Union 
    5 
    6from pydantic_core import CoreSchema, SchemaSerializer, SchemaValidator 
    7 
    8from ..errors import PydanticErrorCodes, PydanticUserError 
    9from ..plugin._schema_validator import PluggableSchemaValidator 
    10 
    11if TYPE_CHECKING: 
    12    from ..dataclasses import PydanticDataclass 
    13    from ..main import BaseModel 
    14    from ..type_adapter import TypeAdapter 
    15 
    16 
    17ValSer = TypeVar('ValSer', bound=Union[SchemaValidator, PluggableSchemaValidator, SchemaSerializer]) 
    18T = TypeVar('T') 
    19 
    20 
    21class MockCoreSchema(Mapping[str, Any]): 
    22    """Mocker for `pydantic_core.CoreSchema` which optionally attempts to 
    23    rebuild the thing it's mocking when one of its methods is accessed and raises an error if that fails. 
    24    """ 
    25 
    26    __slots__ = '_error_message', '_code', '_attempt_rebuild', '_built_memo' 
    27 
    28    def __init__( 
    29        self, 
    30        error_message: str, 
    31        *, 
    32        code: PydanticErrorCodes, 
    33        attempt_rebuild: Callable[[], CoreSchema | None] | None = None, 
    34    ) -> None: 
    35        self._error_message = error_message 
    36        self._code: PydanticErrorCodes = code 
    37        self._attempt_rebuild = attempt_rebuild 
    38        self._built_memo: CoreSchema | None = None 
    39 
    40    def __getitem__(self, key: str) -> Any: 
    41        return self._get_built().__getitem__(key) 
    42 
    43    def __len__(self) -> int: 
    44        return self._get_built().__len__() 
    45 
    46    def __iter__(self) -> Iterator[str]: 
    47        return self._get_built().__iter__() 
    48 
    49    def _get_built(self) -> CoreSchema: 
    50        if self._built_memo is not None: 
    51            return self._built_memo 
    52 
    53        if self._attempt_rebuild: 
    54            schema = self._attempt_rebuild() 
    55            if schema is not None: 
    56                self._built_memo = schema 
    57                return schema 
    58        raise PydanticUserError(self._error_message, code=self._code) 
    59 
    60    def rebuild(self) -> CoreSchema | None: 
    61        self._built_memo = None 
    62        if self._attempt_rebuild: 
    63            schema = self._attempt_rebuild() 
    64            if schema is not None: 
    65                return schema 
    66            else: 
    67                raise PydanticUserError(self._error_message, code=self._code) 
    68        return None 
    69 
    70 
    71class MockValSer(Generic[ValSer]): 
    72    """Mocker for `pydantic_core.SchemaValidator` or `pydantic_core.SchemaSerializer` which optionally attempts to 
    73    rebuild the thing it's mocking when one of its methods is accessed and raises an error if that fails. 
    74    """ 
    75 
    76    __slots__ = '_error_message', '_code', '_val_or_ser', '_attempt_rebuild' 
    77 
    78    def __init__( 
    79        self, 
    80        error_message: str, 
    81        *, 
    82        code: PydanticErrorCodes, 
    83        val_or_ser: Literal['validator', 'serializer'], 
    84        attempt_rebuild: Callable[[], ValSer | None] | None = None, 
    85    ) -> None: 
    86        self._error_message = error_message 
    87        self._val_or_ser = SchemaValidator if val_or_ser == 'validator' else SchemaSerializer 
    88        self._code: PydanticErrorCodes = code 
    89        self._attempt_rebuild = attempt_rebuild 
    90 
    91    def __getattr__(self, item: str) -> None: 
    92        __tracebackhide__ = True 
    93        if self._attempt_rebuild: 
    94            val_ser = self._attempt_rebuild() 
    95            if val_ser is not None: 
    96                return getattr(val_ser, item) 
    97 
    98        # raise an AttributeError if `item` doesn't exist 
    99        getattr(self._val_or_ser, item) 
    100        raise PydanticUserError(self._error_message, code=self._code) 
    101 
    102    def rebuild(self) -> ValSer | None: 
    103        if self._attempt_rebuild: 
    104            val_ser = self._attempt_rebuild() 
    105            if val_ser is not None: 
    106                return val_ser 
    107            else: 
    108                raise PydanticUserError(self._error_message, code=self._code) 
    109        return None 
    110 
    111 
    112def set_type_adapter_mocks(adapter: TypeAdapter) -> None: 
    113    """Set `core_schema`, `validator` and `serializer` to mock core types on a type adapter instance. 
    114 
    115    Args: 
    116        adapter: The type adapter instance to set the mocks on 
    117    """ 
    118    type_repr = str(adapter._type) 
    119    undefined_type_error_message = ( 
    120        f'`TypeAdapter[{type_repr}]` is not fully defined; you should define `{type_repr}` and all referenced types,' 
    121        f' then call `.rebuild()` on the instance.' 
    122    ) 
    123 
    124    def attempt_rebuild_fn(attr_fn: Callable[[TypeAdapter], T]) -> Callable[[], T | None]: 
    125        def handler() -> T | None: 
    126            if adapter.rebuild(raise_errors=False, _parent_namespace_depth=5) is not False: 
    127                return attr_fn(adapter) 
    128            return None 
    129 
    130        return handler 
    131 
    132    adapter.core_schema = MockCoreSchema(  # pyright: ignore[reportAttributeAccessIssue] 
    133        undefined_type_error_message, 
    134        code='class-not-fully-defined', 
    135        attempt_rebuild=attempt_rebuild_fn(lambda ta: ta.core_schema), 
    136    ) 
    137    adapter.validator = MockValSer(  # pyright: ignore[reportAttributeAccessIssue] 
    138        undefined_type_error_message, 
    139        code='class-not-fully-defined', 
    140        val_or_ser='validator', 
    141        attempt_rebuild=attempt_rebuild_fn(lambda ta: ta.validator), 
    142    ) 
    143    adapter.serializer = MockValSer(  # pyright: ignore[reportAttributeAccessIssue] 
    144        undefined_type_error_message, 
    145        code='class-not-fully-defined', 
    146        val_or_ser='serializer', 
    147        attempt_rebuild=attempt_rebuild_fn(lambda ta: ta.serializer), 
    148    ) 
    149 
    150 
    151def set_model_mocks(cls: type[BaseModel], undefined_name: str = 'all referenced types') -> None: 
    152    """Set `__pydantic_core_schema__`, `__pydantic_validator__` and `__pydantic_serializer__` to mock core types on a model. 
    153 
    154    Args: 
    155        cls: The model class to set the mocks on 
    156        undefined_name: Name of the undefined thing, used in error messages 
    157    """ 
    158    undefined_type_error_message = ( 
    159        f'`{cls.__name__}` is not fully defined; you should define {undefined_name},' 
    160        f' then call `{cls.__name__}.model_rebuild()`.' 
    161    ) 
    162 
    163    def attempt_rebuild_fn(attr_fn: Callable[[type[BaseModel]], T]) -> Callable[[], T | None]: 
    164        def handler() -> T | None: 
    165            if cls.model_rebuild(raise_errors=False, _parent_namespace_depth=5) is not False: 
    166                return attr_fn(cls) 
    167            return None 
    168 
    169        return handler 
    170 
    171    cls.__pydantic_core_schema__ = MockCoreSchema(  # pyright: ignore[reportAttributeAccessIssue] 
    172        undefined_type_error_message, 
    173        code='class-not-fully-defined', 
    174        attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_core_schema__), 
    175    ) 
    176    cls.__pydantic_validator__ = MockValSer(  # pyright: ignore[reportAttributeAccessIssue] 
    177        undefined_type_error_message, 
    178        code='class-not-fully-defined', 
    179        val_or_ser='validator', 
    180        attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_validator__), 
    181    ) 
    182    cls.__pydantic_serializer__ = MockValSer(  # pyright: ignore[reportAttributeAccessIssue] 
    183        undefined_type_error_message, 
    184        code='class-not-fully-defined', 
    185        val_or_ser='serializer', 
    186        attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_serializer__), 
    187    ) 
    188 
    189 
    190def set_dataclass_mocks(cls: type[PydanticDataclass], undefined_name: str = 'all referenced types') -> None: 
    191    """Set `__pydantic_validator__` and `__pydantic_serializer__` to `MockValSer`s on a dataclass. 
    192 
    193    Args: 
    194        cls: The model class to set the mocks on 
    195        undefined_name: Name of the undefined thing, used in error messages 
    196    """ 
    197    from ..dataclasses import rebuild_dataclass 
    198 
    199    undefined_type_error_message = ( 
    200        f'`{cls.__name__}` is not fully defined; you should define {undefined_name},' 
    201        f' then call `pydantic.dataclasses.rebuild_dataclass({cls.__name__})`.' 
    202    ) 
    203 
    204    def attempt_rebuild_fn(attr_fn: Callable[[type[PydanticDataclass]], T]) -> Callable[[], T | None]: 
    205        def handler() -> T | None: 
    206            if rebuild_dataclass(cls, raise_errors=False, _parent_namespace_depth=5) is not False: 
    207                return attr_fn(cls) 
    208            return None 
    209 
    210        return handler 
    211 
    212    cls.__pydantic_core_schema__ = MockCoreSchema(  # pyright: ignore[reportAttributeAccessIssue] 
    213        undefined_type_error_message, 
    214        code='class-not-fully-defined', 
    215        attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_core_schema__), 
    216    ) 
    217    cls.__pydantic_validator__ = MockValSer(  # pyright: ignore[reportAttributeAccessIssue] 
    218        undefined_type_error_message, 
    219        code='class-not-fully-defined', 
    220        val_or_ser='validator', 
    221        attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_validator__), 
    222    ) 
    223    cls.__pydantic_serializer__ = MockValSer(  # pyright: ignore[reportAttributeAccessIssue] 
    224        undefined_type_error_message, 
    225        code='class-not-fully-defined', 
    226        val_or_ser='serializer', 
    227        attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_serializer__), 
    228    )