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