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