|
| 1 | +"""Keep `inspect.signature()` accurate for `defer_build` models before their first use. |
| 2 | +
|
| 3 | +The SDK's model bases set `defer_build=True`, so pydantic builds a model's core |
| 4 | +schema, validator and serializer on first use instead of while the class |
| 5 | +statement runs; a module full of models then costs almost nothing to import in |
| 6 | +a process that never validates them. |
| 7 | +
|
| 8 | +The one observable pydantic ties to that build is the class's `__signature__` |
| 9 | +(the synthesized `__init__` signature that `inspect.signature(Model)` reports): |
| 10 | +under `defer_build` it exists only once the model has built, and until then a |
| 11 | +class reports the generic `(**data)`. `DeferredSignature` is a class-level |
| 12 | +`__signature__` that completes the (one-time) build via the public |
| 13 | +`model_rebuild()` on first signature access, so runtime introspection matches an |
| 14 | +eagerly-built model while nothing is built at import. |
| 15 | +""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +from inspect import Signature |
| 20 | +from typing import Any, TypeVar, cast |
| 21 | + |
| 22 | +from pydantic import BaseModel |
| 23 | + |
| 24 | +__all__ = ["DeferredSignature", "deferred_model", "install_deferred_signature", "new_deferred_signature"] |
| 25 | + |
| 26 | +_ModelT = TypeVar("_ModelT", bound=BaseModel) |
| 27 | + |
| 28 | + |
| 29 | +class DeferredSignature: |
| 30 | + """Class-level `__signature__` for a `defer_build=True` pydantic model. |
| 31 | +
|
| 32 | + First access completes the model build via `model_rebuild()` (public API, |
| 33 | + a no-op once complete); pydantic then binds its own signature on the class, |
| 34 | + replacing this attribute, and that is what gets returned. If the build cannot |
| 35 | + complete (an unresolvable forward reference), no signature results and |
| 36 | + `inspect` falls back to `__init__` exactly as for any unbuilt deferred model. |
| 37 | + """ |
| 38 | + |
| 39 | + def __get__(self, obj: object | None, owner: type[BaseModel] | None = None) -> Any: |
| 40 | + if obj is not None or owner is None: |
| 41 | + # Instances take their signature from `__call__`; only the class has one. |
| 42 | + raise AttributeError("__signature__") |
| 43 | + # `raise_errors=False`: an annotation that cannot resolve yet leaves the |
| 44 | + # model unbuilt (and this attribute in place) instead of raising here, |
| 45 | + # so `inspect` falls back to `__init__` like it does for pydantic itself. |
| 46 | + owner.model_rebuild(raise_errors=False) |
| 47 | + bound = vars(owner).get("__signature__") |
| 48 | + if bound is None or bound is self: |
| 49 | + raise AttributeError("__signature__") |
| 50 | + return bound.__get__(None, owner) if hasattr(bound, "__get__") else bound |
| 51 | + |
| 52 | + |
| 53 | +def new_deferred_signature() -> Signature: |
| 54 | + """A `DeferredSignature` typed as the `Signature` it stands in for. |
| 55 | +
|
| 56 | + Assign it in a model class body WITHOUT an annotation (`__signature__ = |
| 57 | + new_deferred_signature()`), matching how `BaseModel` itself binds |
| 58 | + `__signature__`, so it never appears in `get_type_hints(Model)`. |
| 59 | + """ |
| 60 | + return cast(Signature, DeferredSignature()) |
| 61 | + |
| 62 | + |
| 63 | +def install_deferred_signature(cls: type[BaseModel]) -> None: |
| 64 | + """Give a still-deferred model class its own lazily-completing `__signature__`. |
| 65 | +
|
| 66 | + Meant for a base's `__pydantic_init_subclass__` hook: every deferred subclass |
| 67 | + needs its OWN `__signature__` entry, otherwise it would inherit an |
| 68 | + already-built parent's signature. A class that pydantic already completed |
| 69 | + (e.g. one that turned `defer_build` off) keeps the real signature it has. |
| 70 | + """ |
| 71 | + if not cls.__pydantic_complete__: |
| 72 | + setattr(cls, "__signature__", new_deferred_signature()) |
| 73 | + |
| 74 | + |
| 75 | +def deferred_model(cls: type[_ModelT]) -> type[_ModelT]: |
| 76 | + """Class decorator for the root of a `defer_build` model hierarchy under `BaseModel`. |
| 77 | +
|
| 78 | + Installs the lazy `__signature__` on `cls` and, via a `__pydantic_init_subclass__` |
| 79 | + hook, on every model that later subclasses it (a subclass inherits the deferred |
| 80 | + build through the config). Models deriving from the SDK bases (`MCPModel`, |
| 81 | + `WireModel`, `WireRootModel`) get this from the base and do not need it. |
| 82 | + """ |
| 83 | + install_deferred_signature(cls) |
| 84 | + |
| 85 | + def __pydantic_init_subclass__(sub: type[BaseModel], **kwargs: Any) -> None: |
| 86 | + install_deferred_signature(sub) |
| 87 | + # Run whatever hook `cls`'s own bases define (BaseModel's is a no-op). |
| 88 | + # (`super()` from a classmethod defined outside the class body needs |
| 89 | + # both arguments; the two-argument form is opaque to the type checker.) |
| 90 | + cast(Any, super(cls, sub)).__pydantic_init_subclass__(**kwargs) |
| 91 | + |
| 92 | + setattr(cls, "__pydantic_init_subclass__", classmethod(__pydantic_init_subclass__)) |
| 93 | + return cls |
0 commit comments