Coverage for src/dynapydantic/subclass_tracking_model.py: 99%
85 statements
« prev ^ index » next coverage.py v7.14.2, created at 2026-07-13 03:00 +0000
1"""Base class for dynamic pydantic models"""
3import dataclasses
4import inspect
5import typing as ty
7import pydantic
8from pydantic import BaseModel, GetCoreSchemaHandler
9from pydantic_core import PydanticCustomError, core_schema
11from .exceptions import ConfigurationError, Error
12from .tracking_group import TrackingGroup
13from .union_mode import UnionRealization
16class SubclassTrackingModel(pydantic.BaseModel):
17 """Subclass-tracking BaseModel
19 This will inject a [`TrackingGroup`][dynapydantic.TrackingGroup] into your
20 class and automate the registration of subclasses.
22 Similar to `BaseModel`, `SubclassTrackingModel` can take arguments in the
23 class declaration. Arguments from `BaseModel` will be forwarded.
24 Additionally, any fields from `TrackingGroup` will be forwarded to the
25 internal `TrackingGroup` instance. The following additional arguments are
26 supported:
28 1. `exclude_from_union`: This flag is intended to be used with descendents
29 of `SubclassTrackingModel`. If `True`, this subclass will be omitted
30 from tracking. The default for this flag is `True` for direct
31 descendents of `SubclassTrackingModel` and `False` otherwise.
32 2. `union_realization`: When the union should be realized. See
33 [`UnionRealization`][dynapydantic.UnionRealization] for more details
34 on the various options. The default is to realize unions at model
35 construction time.
36 """
38 def __init_subclass__(cls, *args, **kwargs) -> None:
39 """Subclass hook"""
40 # Intercept any kwargs that are intended for TrackingGroup or
41 # __pydantic_init_subclass__
42 sig = inspect.signature(SubclassTrackingModel.__pydantic_init_subclass__)
43 super().__init_subclass__(
44 *args,
45 **{
46 k: v
47 for k, v in kwargs.items()
48 if k not in TrackingGroup.model_fields and k not in sig.parameters
49 },
50 )
52 @classmethod
53 def __pydantic_init_subclass__(
54 cls,
55 *args,
56 exclude_from_union: bool | None = None,
57 union_realization: str | UnionRealization | None = None,
58 **kwargs,
59 ) -> None:
60 """Pydantic subclass hook"""
61 # Forward along any unexpected arguments that were not intended
62 # for TrackingGroup.
63 super().__pydantic_init_subclass__(
64 *args,
65 **{k: v for k, v in kwargs.items() if k not in TrackingGroup.model_fields},
66 )
68 # Initialize the tracking group
69 cls.__DYNAPYDANTIC__: ty.ClassVar[TrackingGroup] = _init_tracking_group(
70 cls, **kwargs
71 )
73 # Initialize our SubclassTrackingModel-specific config
74 cls.__DYNAPYDANTIC_STM_CONFIG__: ty.ClassVar[_StmConfig] = _StmConfig.create(
75 cls,
76 exclude_from_union=exclude_from_union,
77 union_realization=union_realization,
78 inherited=getattr(cls, "__DYNAPYDANTIC_STM_CONFIG__", None),
79 )
81 # If we are going to be tracked, walk the entire MRO (to support
82 # multi-level tree) and register ourselves with each one.
83 if not cls.__DYNAPYDANTIC_STM_CONFIG__.exclude_from_union:
84 for base in cls.__mro__:
85 if (
86 issubclass(base, SubclassTrackingModel)
87 and base is not SubclassTrackingModel
88 and not _is_uninstantiated_generic(base)
89 ):
90 base.__DYNAPYDANTIC__.register_model(cls)
93def _init_tracking_group(
94 cls: type[SubclassTrackingModel],
95 **kwargs,
96) -> TrackingGroup:
97 """Initialize the tracking model embedded in this model"""
98 # If the user already defined one, use it
99 if isinstance((tc := getattr(cls, "tracking_config", None)), TrackingGroup):
100 return tc
102 # Otherwise, we need to make it. We can inherit arguments from our
103 # parent class(es) if they have TrackingGroup's and then allow any
104 # kwargs directly passed here to override.
105 if isinstance(parent_tg := getattr(cls, "__DYNAPYDANTIC__", None), TrackingGroup):
106 tg_kwargs = parent_tg.model_dump(
107 exclude={
108 "name",
109 "models",
110 "discriminator_field",
111 "discriminator_value_generator",
112 }
113 )
114 tg_kwargs |= kwargs
115 if "discriminator_field" in kwargs:
116 tg_kwargs.pop("union_mode", None)
117 else:
118 tg_kwargs = kwargs
119 tg_kwargs.setdefault("name", f"{cls.__name__}-subclasses")
121 try:
122 return TrackingGroup(**tg_kwargs)
123 except pydantic.ValidationError as e:
124 msg = (
125 "SubclassTrackingModel subclasses must either have a "
126 "tracking_config: ClassVar[dynapydantic.TrackingGroup] "
127 "member or pass kwargs sufficient to construct a "
128 "dynapydantic.TrackingGroup in the class declaration. "
129 "The latter approach produced the following "
130 f"ValidationError:\n{e}"
131 )
132 raise ConfigurationError(msg) from e
135@dataclasses.dataclass(frozen=True)
136class _StmConfig:
137 """Config for SubclassTrackingModel"""
139 union_realization: UnionRealization
140 exclude_from_union: bool
142 @classmethod
143 def create(
144 cls,
145 model_t: type[SubclassTrackingModel],
146 *,
147 exclude_from_union: bool | None,
148 union_realization: str | UnionRealization | None = None,
149 inherited: "_StmConfig | None" = None,
150 ) -> "_StmConfig":
151 """Create this model from the user's specified keyword arguments"""
152 # Figure out the union realization time. Prefer direct argument, then
153 # inherited value, then default of model construction time.
154 if union_realization is None:
155 union_realization = (
156 inherited.union_realization
157 if inherited is not None
158 else UnionRealization.MODEL_CONSTRUCTION
159 )
160 elif not isinstance(union_realization, UnionRealization):
161 try:
162 union_realization = UnionRealization(union_realization)
163 except (ValueError, TypeError) as e:
164 msg = f"invalid union_realization: {e}"
165 raise ConfigurationError(msg) from e
167 if exclude_from_union is None:
168 exclude_from_union = _exclude_from_union_default(model_t)
170 return cls(
171 union_realization=union_realization,
172 exclude_from_union=exclude_from_union,
173 )
176def _exclude_from_union_default(model_t: type[SubclassTrackingModel]) -> bool:
177 """Determine the default value for exclude_from_union"""
178 # In general, this shall default to False. It will default to True if:
179 # 1. We are direct descendent of SubclassTrackingModel. This is
180 # because direct descendents tend to be the abstract base classes.
181 if SubclassTrackingModel in model_t.__bases__:
182 return True
184 # 2. We are a generic class with a TypeVar argument (non-concrete).
185 if _is_uninstantiated_generic(model_t):
186 return True
188 # 3. We are a concrete generic class and our origin is a direct
189 # descendent of SubclassTrackingModel. Combined case of 1 and 2. A
190 # concrete generic that is not a direct descendent is the same as any
191 # other class in the middle of an inheritance tree.
192 generic_origin = model_t.__pydantic_generic_metadata__["origin"]
193 if generic_origin is None:
194 return False
195 return SubclassTrackingModel in generic_origin.__bases__
198def _is_uninstantiated_generic(model_t: type[SubclassTrackingModel]) -> bool:
199 """Determine if this a generic model with uninstantiated args"""
200 generic_args = model_t.__pydantic_generic_metadata__["parameters"]
201 return any(isinstance(arg, ty.TypeVar) for arg in generic_args)
204_UNSET = object()
207class ValidationTimeAdapter:
208 """Pydantic type adapter for a dynapydantic-tracked field
210 This adapter returns a validator that evaluates the union at validation time
211 """
213 @staticmethod
214 def __get_pydantic_core_schema__(
215 source_type: type[SubclassTrackingModel],
216 _handler: GetCoreSchemaHandler,
217 ) -> core_schema.CoreSchema:
218 """Get the pydantic schema for this type"""
220 def _validate(value: ty.Any) -> ty.Any: # noqa: ANN401
221 try:
222 adapter = source_type.__DYNAPYDANTIC__.type_adapter
223 except Error as e:
224 err_t = "dynapydantic_error"
225 raise PydanticCustomError(err_t, "{e}", {"e": str(e)}) from e
226 return adapter.validate_python(value)
228 def _serialize(
229 value: BaseModel,
230 info: core_schema.SerializationInfo,
231 ) -> dict[str, ty.Any]:
232 # These arguments we're going to attempt but not require
233 soft_args = (
234 # These were added after 2.0 (we pin >= 2)
235 "context",
236 "exclude_computed_fields",
237 "serialize_as_any",
238 "polymorphic_serialization",
239 )
240 args: dict[str, ty.Any] = {
241 # SerializationInfo doesn't expose warnings, so we have to
242 # pick one option
243 "warnings": False,
244 }
245 for arg in soft_args:
246 if (v := getattr(info, arg, _UNSET)) is not _UNSET: 246 ↛ 245line 246 didn't jump to line 245 because the condition on line 246 was always true
247 args[arg] = v
249 return value.model_dump(
250 mode=info.mode,
251 # Pydantic's types on SerializationInfo's include/exclude don't
252 # match up with the corresponding parameter types on model_dump
253 include=info.include, # type: ignore[bad-argument-type]
254 exclude=info.exclude, # type: ignore[bad-argument-type]
255 by_alias=info.by_alias,
256 exclude_unset=info.exclude_unset,
257 exclude_defaults=info.exclude_defaults,
258 exclude_none=info.exclude_none,
259 round_trip=info.round_trip,
260 **args,
261 )
263 return core_schema.no_info_plain_validator_function(
264 _validate,
265 serialization=core_schema.plain_serializer_function_ser_schema(
266 _serialize,
267 info_arg=True,
268 when_used="unless-none",
269 return_schema=core_schema.dict_schema(
270 core_schema.str_schema(), core_schema.any_schema()
271 ),
272 ),
273 )