Coverage for src/dynapydantic/subclass_tracking_model.py: 99%
130 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 dataclasses
4import functools
5import json
6import typing as ty
8import pydantic
9from pydantic import (
10 BaseModel,
11 GetCoreSchemaHandler,
12 GetJsonSchemaHandler,
13 PydanticInvalidForJsonSchema,
14)
15from pydantic.json_schema import JsonSchemaValue
16from pydantic_core import PydanticCustomError, core_schema
18from .exceptions import ConfigurationError, Error
19from .tracking_group import TrackingGroup
20from .union_mode import UnionRealization
21from .version_check import pydantic_ge
23# Names of kwargs to __pydantic_init_subclass__. This has to be manually kept
24# in sync to avoid a call to inspect.signature on the registration hot path.
25_STM_INIT_SUBCLASS_KWARGS = ("exclude_from_union", "union_realization")
28class SubclassTrackingModel(pydantic.BaseModel):
29 """Subclass-tracking BaseModel
31 This will inject a [`TrackingGroup`][dynapydantic.TrackingGroup] into your
32 class and automate the registration of subclasses.
34 Similar to `BaseModel`, `SubclassTrackingModel` can take arguments in the
35 class declaration. Arguments from `BaseModel` will be forwarded.
36 Additionally, any fields from `TrackingGroup` will be forwarded to the
37 internal `TrackingGroup` instance. The following additional arguments are
38 supported:
40 1. `exclude_from_union`: This flag is intended to be used with descendents
41 of `SubclassTrackingModel`. If `True`, this subclass will be omitted
42 from tracking. The default for this flag is `True` for direct
43 descendents of `SubclassTrackingModel` and `False` otherwise.
44 2. `union_realization`: When the union should be realized. See
45 [`UnionRealization`][dynapydantic.UnionRealization] for more details
46 on the various options. The default is to realize unions at model
47 construction time.
48 """
50 def __init_subclass__(cls, *args, **kwargs) -> None:
51 """Subclass hook"""
52 # Intercept any kwargs that are intended for TrackingGroup or
53 # __pydantic_init_subclass__
54 super().__init_subclass__(
55 *args,
56 **{
57 k: v
58 for k, v in kwargs.items()
59 if k not in TrackingGroup.model_fields
60 and k not in _STM_INIT_SUBCLASS_KWARGS
61 },
62 )
64 @classmethod
65 def __pydantic_init_subclass__(
66 cls,
67 *args,
68 exclude_from_union: bool | None = None,
69 union_realization: str | UnionRealization | None = None,
70 **kwargs,
71 ) -> None:
72 """Pydantic subclass hook"""
73 # Forward along any unexpected arguments that were not intended
74 # for TrackingGroup.
75 super().__pydantic_init_subclass__(
76 *args,
77 **{k: v for k, v in kwargs.items() if k not in TrackingGroup.model_fields},
78 )
80 # Initialize the tracking group
81 cls.__DYNAPYDANTIC__: ty.ClassVar[TrackingGroup] = _init_tracking_group(
82 cls, **kwargs
83 )
85 # Initialize our SubclassTrackingModel-specific config
86 cls.__DYNAPYDANTIC_STM_CONFIG__: ty.ClassVar[_StmConfig] = _StmConfig.create(
87 cls,
88 exclude_from_union=exclude_from_union,
89 union_realization=union_realization,
90 inherited=getattr(cls, "__DYNAPYDANTIC_STM_CONFIG__", None),
91 )
93 # If we are going to be tracked, walk the entire MRO (to support
94 # multi-level tree) and register ourselves with each one.
95 if not cls.__DYNAPYDANTIC_STM_CONFIG__.exclude_from_union:
96 for base in cls.__mro__:
97 if (
98 issubclass(base, SubclassTrackingModel)
99 and base is not SubclassTrackingModel
100 and not _is_uninstantiated_generic(base)
101 ):
102 base.__DYNAPYDANTIC__.register_model(cls)
105def _init_tracking_group(
106 cls: type[SubclassTrackingModel],
107 **kwargs,
108) -> TrackingGroup:
109 """Initialize the tracking model embedded in this model"""
110 # If the user already defined one, use it
111 if isinstance((tc := getattr(cls, "tracking_config", None)), TrackingGroup):
112 return tc
114 # Otherwise, we need to make it. We can inherit arguments from our
115 # parent class(es) if they have TrackingGroup's and then allow any
116 # kwargs directly passed here to override.
117 if isinstance(parent_tg := getattr(cls, "__DYNAPYDANTIC__", None), TrackingGroup):
118 tg_kwargs: dict[str, ty.Any] = {
119 "union_mode": parent_tg.union_mode,
120 "plugin_entry_point": parent_tg.plugin_entry_point,
121 } | kwargs
122 if "discriminator_field" in kwargs:
123 tg_kwargs.pop("union_mode", None)
124 else:
125 tg_kwargs = kwargs
126 tg_kwargs.setdefault("name", f"{cls.__name__}-subclasses")
128 try:
129 return TrackingGroup(**tg_kwargs)
130 except pydantic.ValidationError as e:
131 msg = (
132 "SubclassTrackingModel subclasses must either have a "
133 "tracking_config: ClassVar[dynapydantic.TrackingGroup] "
134 "member or pass kwargs sufficient to construct a "
135 "dynapydantic.TrackingGroup in the class declaration. "
136 "The latter approach produced the following "
137 f"ValidationError:\n{e}"
138 )
139 raise ConfigurationError(msg) from e
142@dataclasses.dataclass(frozen=True)
143class _StmConfig:
144 """Config for SubclassTrackingModel"""
146 union_realization: UnionRealization
147 exclude_from_union: bool
149 @classmethod
150 def create(
151 cls,
152 model_t: type[SubclassTrackingModel],
153 *,
154 exclude_from_union: bool | None,
155 union_realization: str | UnionRealization | None = None,
156 inherited: "_StmConfig | None" = None,
157 ) -> "_StmConfig":
158 """Create this model from the user's specified keyword arguments"""
159 # Figure out the union realization time. Prefer direct argument, then
160 # inherited value, then default of model construction time.
161 if union_realization is None:
162 union_realization = (
163 inherited.union_realization
164 if inherited is not None
165 else UnionRealization.MODEL_CONSTRUCTION
166 )
167 elif not isinstance(union_realization, UnionRealization):
168 try:
169 union_realization = UnionRealization(union_realization)
170 except (ValueError, TypeError) as e:
171 msg = f"invalid union_realization: {e}"
172 raise ConfigurationError(msg) from e
174 if exclude_from_union is None:
175 exclude_from_union = _exclude_from_union_default(model_t)
177 return cls(
178 union_realization=union_realization,
179 exclude_from_union=exclude_from_union,
180 )
183def _exclude_from_union_default(model_t: type[SubclassTrackingModel]) -> bool:
184 """Determine the default value for exclude_from_union"""
185 # In general, this shall default to False. It will default to True if:
186 # 1. We are direct descendent of SubclassTrackingModel. This is
187 # because direct descendents tend to be the abstract base classes.
188 if SubclassTrackingModel in model_t.__bases__:
189 return True
191 # 2. We are a generic class with a TypeVar argument (non-concrete).
192 if _is_uninstantiated_generic(model_t):
193 return True
195 # 3. We are a concrete generic class and our origin is a direct
196 # descendent of SubclassTrackingModel. Combined case of 1 and 2. A
197 # concrete generic that is not a direct descendent is the same as any
198 # other class in the middle of an inheritance tree.
199 generic_origin = model_t.__pydantic_generic_metadata__["origin"]
200 if generic_origin is None:
201 return False
202 return SubclassTrackingModel in generic_origin.__bases__
205def _is_uninstantiated_generic(model_t: type[SubclassTrackingModel]) -> bool:
206 """Determine if this a generic model with uninstantiated args"""
207 generic_args = model_t.__pydantic_generic_metadata__["parameters"]
208 return any(isinstance(arg, ty.TypeVar) for arg in generic_args)
211_UNSET = object()
214class ValidationTimeAdapter:
215 """Pydantic type adapter for a dynapydantic-tracked field
217 This adapter returns a validator that evaluates the union at validation time
218 """
220 @staticmethod
221 def __get_pydantic_core_schema__(
222 source_type: type[SubclassTrackingModel],
223 _handler: GetCoreSchemaHandler,
224 ) -> core_schema.CoreSchema:
225 """Get the pydantic schema for this type"""
227 def _validate(
228 value: ty.Any, # noqa: ANN401
229 info: core_schema.ValidationInfo,
230 *,
231 strict: bool,
232 ) -> ty.Any: # noqa: ANN401
233 try:
234 adapter = source_type.__DYNAPYDANTIC__.type_adapter
235 except Error as e:
236 err_t = "dynapydantic_error"
237 raise PydanticCustomError(err_t, "{e}", {"e": str(e)}) from e
239 kwargs = _validation_kwargs(info)
240 kwargs["strict"] = strict
241 if info.mode == "json":
242 # Field validators receive JSON after the enclosing document has
243 # already been decoded. Re-encode the field so the nested
244 # adapter can apply JSON-specific strict-validation behavior.
245 # https://github.com/pydantic/pydantic/issues/11154
246 kwargs.pop("from_attributes", None)
247 try:
248 value_j = json.dumps(value)
249 # Since the object came from a JSON load, this shouldn't ever
250 # occur, but just being overly defensive.
251 except (TypeError, ValueError, OverflowError) as e:
252 err_t = "json_reencode_failure"
253 msg = "JSON re-encoding failed: {e}"
254 raise PydanticCustomError(err_t, msg, {"e": str(e)}) from e
256 return adapter.validate_json(value_j, **kwargs)
257 return adapter.validate_python(value, **kwargs)
259 def _serialize(
260 value: BaseModel,
261 info: core_schema.SerializationInfo,
262 ) -> dict[str, ty.Any]:
263 # These arguments we're going to attempt but not require
264 soft_args = (
265 # These were added after 2.0 (we pin >= 2)
266 "context",
267 "exclude_computed_fields",
268 "serialize_as_any",
269 "polymorphic_serialization",
270 )
271 args: dict[str, ty.Any] = {
272 # SerializationInfo doesn't expose warnings, so we have to
273 # pick one option
274 "warnings": False,
275 }
276 for arg in soft_args:
277 if (v := getattr(info, arg, _UNSET)) is not _UNSET: 277 ↛ 276line 277 didn't jump to line 276 because the condition on line 277 was always true
278 args[arg] = v
280 return value.model_dump(
281 mode=info.mode,
282 # Pydantic's types on SerializationInfo's include/exclude don't
283 # match up with the corresponding parameter types on model_dump
284 include=info.include, # type: ignore[bad-argument-type]
285 exclude=info.exclude, # type: ignore[bad-argument-type]
286 by_alias=info.by_alias,
287 exclude_unset=info.exclude_unset,
288 exclude_defaults=info.exclude_defaults,
289 exclude_none=info.exclude_none,
290 round_trip=info.round_trip,
291 **args,
292 )
294 _validate_lax = functools.partial(_validate, strict=False)
295 _validate_strict = functools.partial(_validate, strict=True)
297 serialization = core_schema.plain_serializer_function_ser_schema(
298 _serialize,
299 info_arg=True,
300 when_used="unless-none",
301 return_schema=core_schema.dict_schema(
302 core_schema.str_schema(), core_schema.any_schema()
303 ),
304 )
305 metadata = {"dynapydantic_source_type": source_type}
306 return core_schema.lax_or_strict_schema(
307 core_schema.with_info_plain_validator_function(_validate_lax),
308 core_schema.with_info_plain_validator_function(_validate_strict),
309 metadata=metadata,
310 serialization=serialization,
311 )
313 @staticmethod
314 def __get_pydantic_json_schema__(
315 schema: core_schema.CoreSchema,
316 handler: GetJsonSchemaHandler,
317 ) -> JsonSchemaValue:
318 """Lazily build the JSON schema from the currently registered subclasses
320 Runs whenever JSON schema generation actually happens (e.g. a call to
321 `model_json_schema()`), not when the core schema was first built.
322 Reflects whatever subclasses are registered with the `TrackingGroup`
323 at the time of the call.
325 Parameters
326 ----------
327 schema
328 Schema for the field type
329 handler
330 JSON schema handler to convert core_schemas to JSON schemas
332 Returns
333 -------
334 JsonSchemaValue
335 JSON schema for the field
337 Raises
338 ------
339 pydantic.errors.PydanticInvalidForJsonSchema
340 If the JSON schema was unable to be generated.
341 """
342 try:
343 source_type = schema["metadata"]["dynapydantic_source_type"]
344 except KeyError as e:
345 msg = "Missing dynapydantic schema metadata."
346 raise PydanticInvalidForJsonSchema(msg) from e
348 try:
349 union_schema = source_type.__DYNAPYDANTIC__.type_adapter.core_schema
350 except Error as e:
351 msg = str(e)
352 raise PydanticInvalidForJsonSchema(msg) from e
354 return handler(union_schema)
357def _validation_kwargs(
358 info: core_schema.ValidationInfo,
359) -> dict[str, ty.Any]:
360 """Extract keyword arguments for TypeAdapter.validate_python from info."""
361 kwargs: dict[str, ty.Any] = {}
363 if (ctx := getattr(info, "context", None)) is not None:
364 kwargs["context"] = ctx
366 if (config := getattr(info, "config", None)) is not None:
367 # .validate() didn't support extra until 2.12
368 if (
369 pydantic_ge((2, 12, 0))
370 and (val := config.get("extra_fields_behavior")) is not None
371 ):
372 kwargs["extra"] = val
374 if pydantic_ge((2, 11, 0)): 374 ↛ 380line 374 didn't jump to line 380 because the condition on line 374 was always true
375 if (val := config.get("validate_by_alias")) is not None:
376 kwargs["by_alias"] = val
377 if (val := config.get("validate_by_name")) is not None:
378 kwargs["by_name"] = val
380 if (val := config.get("from_attributes")) is not None:
381 kwargs["from_attributes"] = val
383 return kwargs