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