Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import sys
4from collections.abc import Callable, Generator, Iterator, Mapping
5from functools import cached_property
6from types import MappingProxyType
7from typing import Any, NamedTuple, TypeAlias, TypeVar
9from typing_extensions import ParamSpec, TypeAliasType, TypeVarTuple
11from ._contextlib import contextmanager
13GlobalsNamespace: TypeAlias = 'dict[str, Any]'
14"""A global namespace.
16In most cases, this is a reference to the `__dict__` attribute of a module.
17This namespace type is expected as the `globals` argument during annotations evaluation.
18"""
20MappingNamespace: TypeAlias = Mapping[str, Any]
21"""Any kind of namespace.
23In most cases, this is a local namespace (e.g. the `__dict__` attribute of a class,
24the [`f_locals`][frame.f_locals] attribute of a frame object, when dealing with types
25defined inside functions).
26This namespace type is expected as the `locals` argument during annotations evaluation.
27"""
29_TypeVarLike: TypeAlias = 'TypeVar | ParamSpec | TypeVarTuple'
32class NamespacesTuple(NamedTuple):
33 """A tuple of globals and locals to be used during annotations evaluation.
35 This datastructure is defined as a named tuple so that it can easily be unpacked:
37 ```python {lint="skip" test="skip"}
38 def eval_type(typ: type[Any], ns: NamespacesTuple) -> None:
39 return eval(typ, *ns)
40 ```
41 """
43 globals: GlobalsNamespace
44 """The namespace to be used as the `globals` argument during annotations evaluation."""
46 locals: MappingNamespace
47 """The namespace to be used as the `locals` argument during annotations evaluation."""
50def get_module_ns_of(obj: Any) -> dict[str, Any]:
51 """Get the namespace of the module where the object is defined.
53 Caution: this function does not return a copy of the module namespace, so the result
54 should not be mutated. The burden of enforcing this is on the caller.
55 """
56 module_name = getattr(obj, '__module__', None)
57 if module_name:
58 try:
59 return sys.modules[module_name].__dict__
60 except KeyError:
61 # happens occasionally, see https://github.com/pydantic/pydantic/issues/2363
62 return {}
63 return {}
66# Note that this class is almost identical to `collections.ChainMap`, but need to enforce
67# immutable mappings here:
68class LazyLocalNamespace(Mapping[str, Any]):
69 """A lazily evaluated mapping, to be used as the `locals` argument during annotations evaluation.
71 While the [`eval`][eval] function expects a mapping as the `locals` argument, it only
72 performs `__getitem__` calls. The [`Mapping`][collections.abc.Mapping] abstract base class
73 is fully implemented only for type checking purposes.
75 Args:
76 *namespaces: The namespaces to consider, in ascending order of priority.
78 Example:
79 ```python {lint="skip" test="skip"}
80 ns = LazyLocalNamespace({'a': 1, 'b': 2}, {'a': 3})
81 ns['a']
82 #> 3
83 ns['b']
84 #> 2
85 ```
86 """
88 def __init__(self, *namespaces: MappingNamespace) -> None:
89 self._namespaces = namespaces
91 @cached_property
92 def data(self) -> dict[str, Any]:
93 # In rare scenarios, a race condition can happen when iterating over class namespaces (i.e. created with `var(SomeClass)`),
94 # if another thread accesses the type annotations for the first time (since 3.14, reading annotations can set attributes
95 # such as `__annotations_cache__` or `__annotate_func__` to the type).
96 return {
97 k: v for ns in self._namespaces for k, v in (ns.copy() if isinstance(ns, MappingProxyType) else ns).items()
98 }
100 def __len__(self) -> int:
101 return len(self.data)
103 def __getitem__(self, key: str) -> Any:
104 return self.data[key]
106 def __contains__(self, key: object) -> bool:
107 return key in self.data
109 def __iter__(self) -> Iterator[str]:
110 return iter(self.data)
113def ns_for_function(obj: Callable[..., Any], parent_namespace: MappingNamespace | None = None) -> NamespacesTuple:
114 """Return the global and local namespaces to be used when evaluating annotations for the provided function.
116 The global namespace will be the `__dict__` attribute of the module the function was defined in.
117 The local namespace will contain the `__type_params__` introduced by PEP 695.
119 Args:
120 obj: The object to use when building namespaces.
121 parent_namespace: Optional namespace to be added with the lowest priority in the local namespace.
122 If the passed function is a method, the `parent_namespace` will be the namespace of the class
123 the method is defined in. Thus, we also fetch type `__type_params__` from there (i.e. the
124 class-scoped type variables).
125 """
126 locals_list: list[MappingNamespace] = []
127 if parent_namespace is not None:
128 locals_list.append(parent_namespace)
130 # Get the `__type_params__` attribute introduced by PEP 695.
131 # Note that the `typing._eval_type` function expects type params to be
132 # passed as a separate argument. However, internally, `_eval_type` calls
133 # `ForwardRef._evaluate` which will merge type params with the localns,
134 # essentially mimicking what we do here.
135 type_params: tuple[_TypeVarLike, ...] = getattr(obj, '__type_params__', ())
136 if parent_namespace is not None:
137 # We also fetch type params from the parent namespace. If present, it probably
138 # means the function was defined in a class. This is to support the following:
139 # https://github.com/python/cpython/issues/124089.
140 type_params += parent_namespace.get('__type_params__', ())
142 locals_list.append({t.__name__: t for t in type_params})
144 # What about short-circuiting to `obj.__globals__`?
145 globalns = get_module_ns_of(obj)
147 return NamespacesTuple(globalns, LazyLocalNamespace(*locals_list))
150class NsResolver:
151 """A class responsible for the namespaces resolving logic for annotations evaluation.
153 This class handles the namespace logic when evaluating annotations mainly for class objects.
155 It holds a stack of classes that are being inspected during the core schema building,
156 and the `types_namespace` property exposes the globals and locals to be used for
157 type annotation evaluation. Additionally -- if no class is present in the stack -- a
158 fallback globals and locals can be provided using the `namespaces_tuple` argument
159 (this is useful when generating a schema for a simple annotation, e.g. when using
160 `TypeAdapter`).
162 The namespace creation logic is unfortunately flawed in some cases, for backwards
163 compatibility reasons and to better support valid edge cases. See the description
164 for the `parent_namespace` argument and the example for more details.
166 Args:
167 namespaces_tuple: The default globals and locals to use if no class is present
168 on the stack. This can be useful when using the `GenerateSchema` class
169 with `TypeAdapter`, where the "type" being analyzed is a simple annotation.
170 parent_namespace: An optional parent namespace that will be added to the locals
171 with the lowest priority. For a given class defined in a function, the locals
172 of this function are usually used as the parent namespace:
174 ```python {lint="skip" test="skip"}
175 from pydantic import BaseModel
177 def func() -> None:
178 SomeType = int
180 class Model(BaseModel):
181 f: 'SomeType'
183 # when collecting fields, an namespace resolver instance will be created
184 # this way:
185 # ns_resolver = NsResolver(parent_namespace={'SomeType': SomeType})
186 ```
188 For backwards compatibility reasons and to support valid edge cases, this parent
189 namespace will be used for *every* type being pushed to the stack. In the future,
190 we might want to be smarter by only doing so when the type being pushed is defined
191 in the same module as the parent namespace.
193 Example:
194 ```python {lint="skip" test="skip"}
195 ns_resolver = NsResolver(
196 parent_namespace={'fallback': 1},
197 )
199 class Sub:
200 m: 'Model'
202 class Model:
203 some_local = 1
204 sub: Sub
206 ns_resolver = NsResolver()
208 # This is roughly what happens when we build a core schema for `Model`:
209 with ns_resolver.push(Model):
210 ns_resolver.types_namespace
211 #> NamespacesTuple({'Sub': Sub}, {'Model': Model, 'some_local': 1})
212 # First thing to notice here, the model being pushed is added to the locals.
213 # Because `NsResolver` is being used during the model definition, it is not
214 # yet added to the globals. This is useful when resolving self-referencing annotations.
216 with ns_resolver.push(Sub):
217 ns_resolver.types_namespace
218 #> NamespacesTuple({'Sub': Sub}, {'Sub': Sub, 'Model': Model})
219 # Second thing to notice: `Sub` is present in both the globals and locals.
220 # This is not an issue, just that as described above, the model being pushed
221 # is added to the locals, but it happens to be present in the globals as well
222 # because it is already defined.
223 # Third thing to notice: `Model` is also added in locals. This is a backwards
224 # compatibility workaround that allows for `Sub` to be able to resolve `'Model'`
225 # correctly (as otherwise models would have to be rebuilt even though this
226 # doesn't look necessary).
227 ```
228 """
230 def __init__(
231 self,
232 namespaces_tuple: NamespacesTuple | None = None,
233 parent_namespace: MappingNamespace | None = None,
234 ) -> None:
235 self._base_ns_tuple = namespaces_tuple or NamespacesTuple({}, {})
236 self._parent_ns = parent_namespace
237 self._types_stack: list[type[Any] | TypeAliasType] = []
239 @cached_property
240 def types_namespace(self) -> NamespacesTuple:
241 """The current global and local namespaces to be used for annotations evaluation."""
242 if not self._types_stack:
243 # TODO: should we merge the parent namespace here?
244 # This is relevant for TypeAdapter, where there are no types on the stack, and we might
245 # need access to the parent_ns. Right now, we sidestep this in `type_adapter.py` by passing
246 # locals to both parent_ns and the base_ns_tuple, but this is a bit hacky.
247 # we might consider something like:
248 # if self._parent_ns is not None:
249 # # Hacky workarounds, see class docstring:
250 # # An optional parent namespace that will be added to the locals with the lowest priority
251 # locals_list: list[MappingNamespace] = [self._parent_ns, self._base_ns_tuple.locals]
252 # return NamespacesTuple(self._base_ns_tuple.globals, LazyLocalNamespace(*locals_list))
253 return self._base_ns_tuple
255 typ = self._types_stack[-1]
257 globalns = get_module_ns_of(typ)
259 locals_list: list[MappingNamespace] = []
260 # Hacky workarounds, see class docstring:
261 # An optional parent namespace that will be added to the locals with the lowest priority
262 if self._parent_ns is not None:
263 locals_list.append(self._parent_ns)
264 if len(self._types_stack) > 1:
265 first_type = self._types_stack[0]
266 locals_list.append({first_type.__name__: first_type})
268 # Adding `__type_params__` *before* `vars(typ)`, as the latter takes priority
269 # (see https://github.com/python/cpython/pull/120272).
270 # TODO `typ.__type_params__` when we drop support for Python 3.11:
271 type_params: tuple[_TypeVarLike, ...] = getattr(typ, '__type_params__', ())
272 if type_params:
273 # Adding `__type_params__` is mostly useful for generic classes defined using
274 # PEP 695 syntax *and* using forward annotations (see the example in
275 # https://github.com/python/cpython/issues/114053). For TypeAliasType instances,
276 # it is way less common, but still required if using a string annotation in the alias
277 # value, e.g. `type A[T] = 'T'` (which is not necessary in most cases).
278 locals_list.append({t.__name__: t for t in type_params})
280 # TypeAliasType instances don't have a `__dict__` attribute, so the check
281 # is necessary:
282 if hasattr(typ, '__dict__'):
283 locals_list.append(vars(typ))
285 # The `len(self._types_stack) > 1` check above prevents this from being added twice:
286 locals_list.append({typ.__name__: typ})
288 return NamespacesTuple(globalns, LazyLocalNamespace(*locals_list))
290 @contextmanager
291 def push(self, typ: type[Any] | TypeAliasType, /) -> Generator[None]:
292 """Push a type to the stack."""
293 self._types_stack.append(typ)
294 # Reset the cached property:
295 self.__dict__.pop('types_namespace', None)
296 try:
297 yield
298 finally:
299 self._types_stack.pop()
300 self.__dict__.pop('types_namespace', None)