Coverage for src/dynapydantic/tracking_group.py: 100%
145 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-09-06 00:44 +0000
1"""Base class for dynamic pydantic models"""
3import contextlib
4import typing as ty
5import warnings
7import pydantic
8import pydantic.fields
9import pydantic_core
11from .exceptions import (
12 AmbiguousDiscriminatorValueError,
13 NoRegisteredTypesError,
14 RegistrationError,
15)
16from .union_mode import DiscriminatedConfig, UnionMode
19def _inject_discriminator_field(
20 cls: type[pydantic.BaseModel],
21 disc_field: str,
22 value: str,
23) -> pydantic.fields.FieldInfo:
24 """Injects the discriminator field into the given model
26 Parameters
27 ----------
28 cls
29 The BaseModel subclass
30 disc_field
31 Name of the discriminator field
32 value
33 Value of the discriminator field
34 """
35 if hasattr(cls, disc_field):
36 msg = (
37 f'Cannot inject discriminator field "{disc_field}" into '
38 f"{cls.__name__}: an attribute with that name already exists. "
39 "Rename either the attribute or the discriminator_field to avoid "
40 "the conflict."
41 )
42 raise RegistrationError(msg)
44 cls.model_fields[disc_field] = pydantic.fields.FieldInfo(
45 default=value,
46 annotation=ty.Literal[value], # type: ignore[not-a-type]
47 frozen=True,
48 )
49 with contextlib.suppress(pydantic.errors.PydanticUndefinedAnnotation):
50 cls.model_rebuild(force=True)
51 return cls.model_fields[disc_field]
54class TrackingGroup(pydantic.BaseModel):
55 """Tracker for pydantic models"""
57 name: str = pydantic.Field(
58 description=(
59 "Name of the tracking group. This is for human display, so it "
60 "doesn't technically need to be globally unique, but it should be "
61 "meaningfully named, as it will be used in error messages."
62 ),
63 )
64 union_mode: UnionMode | None = pydantic.Field(
65 None,
66 description=(
67 "Union validation strategy. Pass a DiscriminatedConfig instance "
68 'or one of the plain strings "smart" or "left_to_right". You can '
69 "also just pass the fields for DiscriminatedConfig to this "
70 "model and they will be forwarded."
71 ),
72 )
73 discriminator_field: str | None = pydantic.Field(
74 None,
75 description=(
76 "Name of the discriminator field. NOTE: This field is "
77 "here as an alias for union_mode.discriminator_field. Passing "
78 "both a discriminator_field and a union_mode will result in an "
79 "error."
80 ),
81 )
82 discriminator_value_generator: ty.Callable[[type], str] | None = pydantic.Field(
83 None,
84 description=(
85 "A callable that produces default values for the discriminator field"
86 ),
87 )
88 plugin_entry_point: str | None = pydantic.Field(
89 None,
90 description=(
91 "If given, then plugins packages will be supported through this "
92 "Python entrypoint. The entrypoint can either be a function, "
93 "which will be called, or simply a module, which will be "
94 "imported. In either case, models found along the import path of "
95 "the entrypoint will be registered. If the entrypoint is a "
96 "function, additional models may be declared in the function."
97 ),
98 )
99 models: dict[str, type[pydantic.BaseModel]] = pydantic.Field(
100 {},
101 description="The tracked models",
102 )
104 _generation: int = pydantic.PrivateAttr(default=0)
105 _adapter: pydantic.TypeAdapter | None = pydantic.PrivateAttr(default=None)
106 _adapter_generation: int = pydantic.PrivateAttr(default=-1)
108 @pydantic.model_validator(mode="after")
109 def _ensure_union_mode(self) -> "TrackingGroup":
110 """There must be a union_mode
112 This validator works as a guard on _coerce_union_mode to make
113 """
114 if self.union_mode is None:
115 msg = (
116 "union_mode is required. This normally indicates that you "
117 "subclasses TrackingGroup and wrote an invalid validator, but "
118 "could also be a bug with dynapydantic, so please file a bug "
119 "report with a reproducer on how you got here if you suspect "
120 "a bug."
121 )
122 raise ValueError(msg)
124 # Ensure the top-level fields are in-sync
125 if isinstance(self.union_mode, DiscriminatedConfig):
126 self.discriminator_field = self.union_mode.discriminator_field
127 self.discriminator_value_generator = (
128 self.union_mode.discriminator_value_generator
129 )
130 else:
131 self.discriminator_field = None
132 self.discriminator_value_generator = None
134 return self
136 @pydantic.model_validator(mode="before")
137 @classmethod
138 def _coerce_union_mode(cls, data: ty.Any) -> ty.Any: # noqa: ANN401
139 """Coerce flat discriminator kwargs into a DiscriminatedConfig.
141 Allows callers to pass ``discriminator_field`` and
142 ``discriminator_value_generator`` at the top level and transparently
143 assembles a ``DiscriminatedConfig`` from them. This avoids an extra
144 import/nesting layer for the user.
145 """
146 if not isinstance(data, dict):
147 return data
149 disc_field = data.get("discriminator_field", None)
150 has_disc_field = disc_field is not None
151 union_mode = data.get("union_mode", None)
152 has_union_mode = union_mode is not None
154 # If the user passed us both a discriminator field and a union_mode,
155 # things must be perfectly consistent
156 if has_disc_field and has_union_mode:
157 consistent = (
158 isinstance(union_mode, DiscriminatedConfig)
159 and disc_field == union_mode.discriminator_field
160 and data.get("discriminator_value_generator")
161 is union_mode.discriminator_value_generator
162 ) or (
163 isinstance(union_mode, dict)
164 and disc_field == union_mode.get("discriminator_field")
165 and data.get("discriminator_value_generator")
166 is union_mode.get("discriminator_value_generator")
167 )
168 if not consistent:
169 msg = (
170 "Received both union_mode and discriminator_field; pass one "
171 "or the other."
172 )
173 raise ValueError(msg)
175 if has_disc_field and not has_union_mode:
176 # Forward arguments to DiscriminatedConfig
177 data["union_mode"] = {
178 "discriminator_field": disc_field,
179 "discriminator_value_generator": data.get(
180 "discriminator_value_generator",
181 ),
182 }
183 elif not has_disc_field and not has_union_mode:
184 msg = "Either union_mode or discriminator_field must be given"
185 raise ValueError(msg)
187 return data
189 @property
190 def _discriminated(self) -> DiscriminatedConfig | None:
191 """Return the DiscriminatedMode config, or None if not discriminated."""
192 return (
193 self.union_mode
194 if isinstance(self.union_mode, DiscriminatedConfig)
195 else None
196 )
198 def load_plugins(self) -> None:
199 """Load plugins to discover/register additional models"""
200 if self.plugin_entry_point is None:
201 return
203 from importlib.metadata import entry_points # noqa: PLC0415
205 for ep in entry_points().select(group=self.plugin_entry_point):
206 plugin = ep.load()
207 if callable(plugin):
208 plugin()
210 @ty.overload
211 def register(self, value: str | None = None) -> ty.Callable[[type], type]: ...
213 @ty.overload
214 def register(self, value: type[pydantic.BaseModel]) -> type[pydantic.BaseModel]: ...
216 def register(
217 self,
218 value: str | type[pydantic.BaseModel] | None = None,
219 ) -> ty.Callable[[type], type] | type[pydantic.BaseModel]:
220 """Register a model into this group (decorator)
222 Parameters
223 ----------
224 value
225 Value for the discriminator field. If not given, then
226 discriminator_value_generator must be non-None or the
227 discriminator field must be declared by hand. Can also be the type
228 itself to register (if the ()'s are omitted from the decorator).
229 """
230 if isinstance(value, type):
231 self.register_model(value)
232 return value
234 def _wrapper(cls: type[pydantic.BaseModel]) -> type[pydantic.BaseModel]:
235 self.register_model(cls, value)
236 return cls
238 return _wrapper
240 def register_model(
241 self,
242 cls: type[pydantic.BaseModel],
243 discriminator_value: str | None = None,
244 ) -> None:
245 """Register the given model into this group
247 Parameters
248 ----------
249 cls
250 The model to register
251 discriminator_value
252 Value for the discriminator field. If not given, then
253 discriminator_value_generator must be non-None or the
254 discriminator field must be declared by hand.
255 """
256 if discriminator_value is not None and not isinstance(discriminator_value, str):
257 msg = (
258 "discriminator_value must be a str if given, was "
259 f"{type(discriminator_value).__name__}"
260 )
261 raise RegistrationError(msg)
263 if not isinstance(cls, type) or not issubclass(cls, pydantic.BaseModel):
264 msg = (
265 "only pydantic BaseModel subclasses can be registered in a "
266 f"TrackingGroup. Got {cls}, which was not."
267 )
268 raise RegistrationError(msg)
270 if (dm := self._discriminated) is not None:
271 disc = dm.discriminator_field
272 field = cls.model_fields.get(disc)
274 if field is None:
275 if discriminator_value is not None:
276 _inject_discriminator_field(cls, disc, discriminator_value)
277 elif dm.discriminator_value_generator is not None:
278 _inject_discriminator_field(
279 cls,
280 disc,
281 dm.discriminator_value_generator(cls),
282 )
283 else:
284 msg = (
285 f"unable to determine a discriminator value for "
286 f'{cls.__name__} in tracking group "{self.name}". '
287 "No value was passed, discriminator_value_generator "
288 f'was None and the "{disc}" field was not defined.'
289 )
290 raise RegistrationError(msg)
291 elif ty.get_origin(field.annotation) is not ty.Literal:
292 msg = (
293 f'the discriminator field "{disc}" already existed in '
294 f"{cls.__name__}, but its type annotation was "
295 f"{field.annotation}, not Literal."
296 )
297 raise RegistrationError(msg)
298 elif (
299 discriminator_value is not None and field.default != discriminator_value
300 ):
301 msg = (
302 f"the discriminator value for {cls.__name__} was "
303 f'ambiguous, the passed value was "{discriminator_value}" '
304 f' and "{field.default}" via the discriminator '
305 f"field ({disc})."
306 )
307 raise AmbiguousDiscriminatorValueError(msg)
309 self._register_with_discriminator_field(cls)
310 else:
311 if discriminator_value is not None:
312 warnings.warn(
313 f'A discriminator_value of "{discriminator_value}" was '
314 f"explicitly passed for {cls.__name__}, but "
315 f'union_mode="{self.union_mode}" does not use a '
316 "discriminator. The value will be ignored.",
317 stacklevel=2,
318 )
319 self._register_plain(cls)
321 def union(
322 self,
323 *,
324 plain: bool | None = None,
325 ) -> ty.Any: # noqa: ANN401
326 """Return the union of all registered models
328 Parameters
329 ----------
330 plain
331 If set to `True`, a plain union of all members will be returned.
332 Otherwise, the returned union will be annotated in accordance with
333 the union mode.
335 Returns
336 -------
337 Any
338 If there is 1 registered type, the type itself. If there is > 1, a
339 union of all registered types. This union may be annotated if
340 `plain` is not `True`.
342 Raises
343 ------
344 NoRegisteredTypesError
345 If no types have been registered yet.
346 """
347 n = len(self.models)
348 if n == 0:
349 msg = (
350 "Unable to produce a union from the tracking group "
351 f'"{self.name}", as no types have been registered yet.'
352 )
353 raise NoRegisteredTypesError(msg)
354 if n == 1:
355 return next(iter(self.models.values()))
357 union_mode = "smart" if plain else self.union_mode
359 if isinstance(union_mode, DiscriminatedConfig):
360 return ty.Annotated[
361 # This is significantly faster than |'ing into a UnionType
362 ty.Union[ # noqa: UP007
363 tuple( # type: ignore[not-a-type]
364 ty.Annotated[x, pydantic.Tag(v)] for v, x in self.models.items()
365 )
366 ],
367 pydantic.Field(discriminator=union_mode.discriminator_field),
368 ]
370 plain_union = ty.Union[ # noqa: UP007
371 tuple(self.models.values()) # type: ignore[not-a-type]
372 ]
373 if union_mode == "left_to_right":
374 return ty.Annotated[plain_union, pydantic.Field(union_mode="left_to_right")]
376 # "smart" mode is pydantic's default behavior on a plain union
377 return plain_union
379 @property
380 def generation(self) -> int:
381 """The generation of the tracking group.
383 This is a counter that increments every time a new registration occurs
384 """
385 return self._generation
387 @property
388 def type_adapter(self) -> pydantic.TypeAdapter:
389 """Get the pydantic TypeAdapter for the union of all group members"""
390 if self.generation != self._adapter_generation:
391 self._adapter = pydantic.TypeAdapter(self.union())
392 self._adapter_generation = self.generation
394 # casting because the if statement ensures it is non-None (because
395 # _adapter_generation starts at -1 and generation increments from 0.
396 return ty.cast("pydantic.TypeAdapter", self._adapter)
398 def _register_with_discriminator_field(self, cls: type[pydantic.BaseModel]) -> None:
399 """Register the model with the default of the discriminator field
401 Parameters
402 ----------
403 cls
404 The class to register, must have the disciminator field set with a
405 unique default value in the group.
406 """
407 disc = ty.cast("DiscriminatedConfig", self.union_mode).discriminator_field
408 value = cls.model_fields[disc].default
409 if value == pydantic_core.PydanticUndefined:
410 msg = (
411 f"{cls.__name__}.{disc} had no default value, it must "
412 "have one which is unique among all tracked models."
413 )
414 raise RegistrationError(msg)
415 if not isinstance(value, str):
416 msg = (
417 f"{cls.__name__}.{disc} had a default value of {value}, which "
418 f"was of type {type(value).__name__}, not str."
419 )
420 raise RegistrationError(msg)
422 self._do_register(value, cls)
424 def _register_plain(self, cls: type[pydantic.BaseModel]) -> None:
425 """Register the model keyed by its class name.
427 Used for smart / left_to_right modes where no discriminator field
428 is involved.
430 Parameters
431 ----------
432 cls
433 The model to register.
434 """
435 self._do_register(str(id(cls)), cls)
437 def _do_register(self, key: str, cls: type[pydantic.BaseModel]) -> None:
438 """Register the given model under the given key
440 Parameters
441 ----------
442 key
443 The key under which to register the model
444 cls
445 The model to register.
446 """
447 if (other := self.models.get(key)) is not None:
448 if other is not cls:
449 msg = (
450 f'Cannot register {cls.__name__} under the "{key}" '
451 f"identifier, which is already in use by {other.__name__}."
452 )
453 raise RegistrationError(msg)
454 else:
455 self._generation += 1
456 self.models[key] = cls