Coverage for core / src / sensorkit / common / model.py: 92%

1# SPDX-License-Identifier: Apache-2.0 

2"""Pydantic model registry infrastructure for discriminated-union validation.""" 

3 

4from __future__ import annotations 

5 

6import collections 

7import contextlib 

8import functools 

9import inspect 

10import warnings 

11from abc import ABC, abstractmethod 

12from collections.abc import Iterable 

13from contextvars import ContextVar 

14from dataclasses import dataclass 

15from types import MappingProxyType 

16from typing import Annotated, Any, ClassVar, Self, override 

17 

18from loguru import logger 

19from pydantic import ( 

20 BaseModel, 

21 GetCoreSchemaHandler, 

22 ModelWrapValidatorHandler, 

23 SerializerFunctionWrapHandler, 

24 TypeAdapter, 

25 ValidationInfo, 

26 WrapSerializer, 

27 WrapValidator, 

28) 

29from pydantic_core import PydanticUndefined 

30 

31 

32class RegistryError(Exception): 

33 """Raised for tag conflicts during ModelRegistryView resolution.""" 

34 

35 

36@dataclass(frozen=True, slots=True) 

37class RegistryEntry: 

38 """Metadata record for a type registered in a ModelRegistry.""" 

39 

40 model_type: type 

41 tag: str | None 

42 namespace: str | None 

43 

44 

45class ModelRegistryBase[T](ABC): 

46 """Base for model registry types.""" 

47 

48 DISCRIMINATOR_CONTEXT: ClassVar[str] = "discriminator" 

49 REGISTRY_CONTEXT: ClassVar[str] = "registry" 

50 

51 @abstractmethod 

52 def validate(self, data: Any, *, info: ValidationInfo | None = None) -> T: 

53 """Validate `data` and return a typed model instance from this registry.""" 

54 ... 

55 

56 @abstractmethod 

57 def __contains__(self, val: type[T]) -> bool: ... 

58 

59 

60def extract_field_default_value(type_: type, field: str | None): 

61 """Return the default value of `field` on `type_`, checking class dict and pydantic field info.""" 

62 if not field: 

63 return None 

64 

65 # Only defaults declared directly on this class, not inherited attrs. 

66 if field in type_.__dict__: 

67 value = type_.__dict__[field] 

68 

69 if value is not None: 

70 return value 

71 

72 if issubclass(type_, BaseModel): 

73 field_info = type_.model_fields.get(field) 

74 

75 if field_info is not None and field_info.default is not PydanticUndefined: 

76 return field_info.default 

77 

78 return None 

79 

80 

81class ModelRegistry[T](ModelRegistryBase[T]): 

82 """A mutable collection of pydantic models. 

83 

84 Types register here at import time via `__init_subclass__` hooks or decorators, with 

85 optional namespace metadata. 

86 

87 Operations are not thread-safe. 

88 """ 

89 

90 def __init__( 

91 self, 

92 *models: type[T], 

93 discriminator: str | None = None, 

94 default_tag: str | None = None, 

95 ): 

96 self.discriminator_field = discriminator 

97 self.default_tag = default_tag 

98 self._entries: dict[type[T], RegistryEntry] = {} 

99 self._tag_index: dict[str, dict[str | None, type[T]]] = collections.defaultdict(dict) 

100 self._update_default_view = False 

101 

102 for model_type in models: 

103 self.add(model_type) 

104 

105 def add(self, model_type: type[T], *, tag: str | None = None, namespace: str | None = None): 

106 """Register a model type in the registry. 

107 

108 The next time `validate()` is called, the registry cache will be updated to include the new 

109 model. This interaction is not thread safe. 

110 

111 Args: 

112 model_type: The model class to register. Abstract classes are ignored and trigger a 

113 warning. 

114 tag: Optional discriminator tag value. If None, will be extracted from the model's 

115 discriminator field default value. 

116 namespace: Optional namespace for organizing models with the same tag. 

117 

118 Raises: 

119 RegistryError: If the model type is already registered, if tag cannot be determined 

120 when no discriminator is set, if a tag conflict occurs in the same namespace, 

121 or if the registry has no discriminator field configured. 

122 """ 

123 if inspect.isabstract(model_type): 

124 warnings.warn( 

125 f"cannot add abstract type {model_type.__name__} to registry", stacklevel=2 

126 ) 

127 return 

128 

129 if tag is None: 

130 if self.discriminator_field is None: 

131 raise RegistryError("cannot add model to registry without a discriminator value") 

132 

133 tag = extract_field_default_value(model_type, self.discriminator_field) 

134 

135 if tag in (None, PydanticUndefined) and self.default_tag is None: 

136 raise RegistryError( 

137 f"could not determine discriminator value for {model_type.__name__}" 

138 ) 

139 

140 if model_type in self._entries: 

141 raise RegistryError(f"model type {model_type.__name__} already registered") 

142 

143 self._entries[model_type] = RegistryEntry( 

144 model_type=model_type, tag=tag, namespace=namespace 

145 ) 

146 

147 if tag is not None: 

148 if namespace in self._tag_index[tag]: 

149 raise RegistryError(f"tag '{tag}' already defined in namespace '{namespace}'") 

150 

151 self._tag_index[tag][namespace] = model_type 

152 

153 self._update_default_view = True 

154 

155 @property 

156 def entries(self): 

157 """A read-only mapping of registered model types to their `RegistryEntry` metadata.""" 

158 return MappingProxyType(self._entries) 

159 

160 def get_tags(self): 

161 """Return all discriminator tags currently registered in this registry.""" 

162 return tuple(self._tag_index.keys()) 

163 

164 def get_namespaces(self, tag: str): 

165 """Get all namespaces for a given tag.""" 

166 if tag in self._tag_index: 

167 return tuple(self._tag_index[tag].keys()) 

168 return () 

169 

170 def get_type(self, tag: str, namespace: str | None = None): 

171 """Return the model type registered under `tag` in the given `namespace`, or `None`.""" 

172 return self._tag_index[tag].get(namespace) 

173 

174 def create_resolved_view(self, namespace_precedence: Iterable[str]) -> ModelRegistryView: 

175 """Create a `ModelRegistryView` of this registry with key conflict resolution.""" 

176 return ModelRegistryView(self, namespace_precedence) 

177 

178 @functools.cached_property 

179 def _default_view(self): 

180 return self.create_resolved_view(namespace_precedence=()) 

181 

182 def discriminator(self): 

183 """Return a `RegistryDiscriminator` pydantic annotation backed by this registry.""" 

184 return RegistryDiscriminator(self) 

185 

186 @functools.cache 

187 def type_adapter(self, model_type: type[T]) -> TypeAdapter: 

188 """Get a cached TypeAdapter for a given model type.""" 

189 return TypeAdapter(model_type) 

190 

191 @override 

192 def validate(self, data: Any, *, info: ValidationInfo | None = None) -> T: 

193 """Validate data against registered models. 

194 

195 This method delegates validation to the default ModelRegistryView, which resolves 

196 model types based on discriminator tags without namespace precedence. The view is 

197 automatically updated if new models have been added to the registry since the last 

198 validation. 

199 

200 Args: 

201 data: The data to validate. Can be a dict, model instance, or any object with 

202 the discriminator field. 

203 info: Optional validation context information that may contain discriminator 

204 tag overrides or registry view references. 

205 

206 Returns: 

207 A validated instance of one of the registered model types. 

208 

209 Raises: 

210 ValueError: If no discriminator tag is available, if no model is found for the 

211 tag, or if there's a discriminator mismatch. 

212 RegistryError: If tag conflicts exist in the default namespace resolution. 

213 """ 

214 if self._update_default_view: 

215 self._default_view.update() 

216 self._update_default_view = False 

217 

218 return self._default_view.validate(data, info=info) 

219 

220 @override 

221 def __contains__(self, val: type[T]) -> bool: 

222 return val in self._entries 

223 

224 

225_registry_var: ContextVar[ModelRegistryView | None] = ContextVar("model_registry", default=None) 

226"""Active ModelRegistryView for the current execution context, if any.""" 

227 

228 

229class ModelRegistryView[T](ModelRegistryBase[T]): 

230 """Immutable view of a ModelRegistry with effective keys resolved by namespace precedence.""" 

231 

232 def __init__(self, registry: ModelRegistry, precedence: Iterable[str] = ()): 

233 self._registry = registry 

234 self._precedence: dict[str, int] = {ns: i for i, ns in enumerate(precedence)} 

235 self._resolved: dict[str, type[T]] = {} 

236 self._index: dict[type[T], str] = {} 

237 

238 self.update() 

239 

240 def _resolve(self, tag: str, namespaces: tuple[str | None, ...]): 

241 logger.debug(f"resolving tag '{tag}' with namespaces: {namespaces}") 

242 best_prio = (None, float("inf")) 

243 

244 for ns in namespaces: 

245 if ns is None: 

246 if best_prio[0] is None: 

247 best_prio = (ns, float("inf")) 

248 else: 

249 prio = self._precedence.get(ns, float("inf")) 

250 

251 if prio < best_prio[1]: 

252 best_prio = (ns, prio) 

253 

254 return best_prio[0] 

255 

256 def update(self): 

257 """Update the registry view based on the current state of the underlying registry.""" 

258 self._resolved.clear() 

259 self._index.clear() 

260 

261 for tag in self._registry.get_tags(): 

262 namespaces = self._registry.get_namespaces(tag) 

263 assert namespaces 

264 

265 # Determine the namespace that takes precedence. 

266 resolved = self._resolve(tag, namespaces) if len(namespaces) > 1 else namespaces[0] 

267 

268 # Build the forward and reverse tag to model mappings. 

269 self._resolved[tag] = self._registry.get_type(tag, namespace=resolved) 

270 self._index[self._resolved[tag]] = tag 

271 

272 def validate(self, data: Any, *, info: ValidationInfo | None = None) -> T: 

273 """Validate data against this registry view.""" 

274 # Determine the tag to use for validation. 

275 tag_from_context: str | None = ( 

276 info.context.get(self.DISCRIMINATOR_CONTEXT) if info and info.context else None 

277 ) 

278 tag_from_data = None 

279 

280 if discriminator := self._registry.discriminator_field: 

281 match data: 

282 case dict(): 

283 tag_from_data = data.get(discriminator) 

284 case _: 

285 tag_from_data = getattr(data, discriminator, None) 

286 

287 # If both data and context provide a discriminator and they differ, error out. 

288 if tag_from_data and tag_from_context and tag_from_data != tag_from_context: 

289 raise ValueError( 

290 f"tag mismatch: data='{tag_from_data}' context='{tag_from_context}'" 

291 ) 

292 

293 model_type: type[T] | None = None 

294 tag: str | None = None 

295 

296 if tag := tag_from_data or tag_from_context: 

297 model_type = self._resolved.get(tag) 

298 

299 if model_type is None: 

300 tag = self._registry.default_tag 

301 model_type = self._resolved.get(tag) 

302 

303 if model_type is None: 

304 raise ValueError(f"no model resolved for tag: {tag}") 

305 

306 if not isinstance(data, dict): 

307 if type(data) is not model_type: 

308 raise ValueError(f"model {model_type.__name__} does not match tag: {tag}") 

309 

310 # Strip the discriminator context, if any. This is necessary to prevent downstream 

311 # registry-backed fields from receiving and interpreting this themselves. 

312 # FIXME: discriminator-from-context should be reviewed, as it may be avoidable by doing 

313 # namespace lookups against the registry, e.g. in `validate_keyword`. This is 

314 # slightly more complicated than it sounds to ensure views/namespaces are folded in. 

315 if tag_from_context: 

316 del info.context[self.DISCRIMINATOR_CONTEXT] 

317 

318 return self._registry.type_adapter(model_type).validate_python(data, context=info.context) 

319 

320 def __contains__(self, val: type[T]) -> bool: 

321 """Return `True` if `val` is a model type resolved by this view.""" 

322 return val in self._index 

323 

324 @contextlib.contextmanager 

325 def as_current(self): 

326 """Context manager that installs this registry as the active ContextVar.""" 

327 token = _registry_var.set(self) 

328 try: 

329 yield self 

330 finally: 

331 _registry_var.reset(token) 

332 

333 

334class RegistryDiscriminator: 

335 """A pydantic discriminator implementation backed by a model registry or view.""" 

336 

337 def __init__(self, registry: ModelRegistry): 

338 self.registry = registry 

339 

340 def _validate(self, data: Any, _handler: ModelWrapValidatorHandler, info: ValidationInfo): 

341 model_type = type(data) 

342 

343 # Look for a registry view in the validation context. If none is found, fallback to the 

344 # contextvar. 

345 registry_view: ModelRegistryBase | None = ( 

346 info.context and info.context.get(ModelRegistryView.REGISTRY_CONTEXT) 

347 ) or _registry_var.get() 

348 

349 # If a concrete model instance is provided, accept it if it's part of the registry. 

350 if isinstance(data, BaseModel): 

351 exists = model_type in registry_view if registry_view else model_type in self.registry 

352 

353 if not exists: 

354 raise ValueError(f"model {model_type.__name__} not in registry") 

355 

356 return data 

357 

358 # Validate using the registry view if available. This is guaranteed to be conflict-free 

359 # since namespaces have been resolved. 

360 if registry_view is not None: 

361 return registry_view.validate(data, info=info) 

362 

363 # Otherwise, validate against the entire deep registry. This will fail if the tag is 

364 # defined in multiple namespaces. 

365 return self.registry.validate(data, info=info) 

366 

367 def _serialize(self, data: Any, handler: SerializerFunctionWrapHandler): 

368 if isinstance(data, BaseModel): 

369 return data.model_dump() 

370 

371 return handler(data) 

372 

373 def __get_pydantic_core_schema__( 

374 self, 

375 source_type: Any, 

376 handler: GetCoreSchemaHandler, 

377 ): 

378 return handler( 

379 Annotated[ 

380 source_type, 

381 WrapValidator(self._validate), 

382 WrapSerializer(self._serialize), 

383 ] 

384 ) 

385 

386 

387class RegistryBaseModel(BaseModel, ABC): 

388 """Base class for pydantic models that auto-register into a ModelRegistry upon subclassing. 

389 

390 This abstract base class does automatic registration of model subclasses into a ModelRegistry 

391 using the discriminator pattern. Each concrete subclass is registered when defined, with its 

392 class name used as the discriminator value by default. 

393 

394 Subclasses must implement get_registry() to specify which ModelRegistry instance they should 

395 register with. The registry's discriminator field is automatically populated on each subclass 

396 with the appropriate discriminator value. 

397 

398 The registration happens in two hooks: 

399 - __init_subclass__: Registers the class in the registry and sets up the discriminator field 

400 - __pydantic_init_subclass__: Ensures discriminator values are available at class scope 

401 

402 Example: 

403 >>> registry = ModelRegistry(discriminator="type") 

404 >>> class MyBase(RegistryBaseModel): 

405 ... type: str 

406 ... 

407 ... @classmethod 

408 ... def get_registry(cls): 

409 ... return registry 

410 >>> 

411 >>> class MyModel(MyBase): 

412 ... pass # Automatically registered with discriminator value "MyModel" 

413 

414 Note: 

415 The root RegistryBaseModel class and direct subclasses that define get_registry() are 

416 not registered themselves - only their concrete descendant classes are registered. 

417 """ 

418 

419 @classmethod 

420 @abstractmethod 

421 def model_registry(cls) -> ModelRegistry[Self]: 

422 """Return the `ModelRegistry` that this model class family registers into.""" 

423 raise NotImplementedError 

424 

425 @classmethod 

426 def model_tag(cls): 

427 """Get the effective discriminator value for this model class. 

428 

429 Raises: 

430 KeyError: the model is not present in the registry 

431 """ 

432 return cls.model_registry().entries[cls].tag 

433 

434 @classmethod 

435 def __init_subclass__(cls, **kwargs: Any): 

436 super().__init_subclass__(**kwargs) 

437 

438 # Skip abstract classes. 

439 if inspect.isabstract(cls): 

440 return 

441 

442 # Skip the root registered model class itself (e.g. Event), only process its subclasses. 

443 if cls is cls._find_registered_base(): 

444 return 

445 

446 registry = cls.model_registry() 

447 field = registry.discriminator_field 

448 tag = extract_field_default_value(cls, field) 

449 

450 if tag is None: 

451 tag = cls.__name__ 

452 

453 if field: 

454 # RegistryBaseModel subclasses are expected to narrow the discriminator field type 

455 # (e.g. `op: Literal["foo"] = "foo"`), which pydantic treats as field shadowing and 

456 # warns about. Suppress that warning for this specific field name. 

457 warnings.filterwarnings( 

458 "ignore", 

459 message=rf'Field name "{field}" in ".*{cls.__name__}" shadows an attribute', 

460 category=UserWarning, 

461 ) 

462 

463 registry.add(cls, tag=tag) 

464 

465 def model_post_init(self, __context: Any): 

466 """Ensure the registry's discriminator field is populated in each instance.""" 

467 registry = self.model_registry() 

468 field = registry.discriminator_field 

469 cls = type(self) 

470 

471 if cls not in registry.entries: 

472 raise RegistryError(f"Model {type(self).__name__} is not registered") 

473 

474 tag = registry.entries[cls].tag 

475 

476 if tag != registry.default_tag: 

477 object.__setattr__(self, field, tag) 

478 

479 @classmethod 

480 def _find_registered_base(cls) -> type | None: 

481 """Find the first class in the MRO that directly inherits from RegisteredBaseModel.""" 

482 for klass in cls.__mro__: 

483 if RegistryBaseModel in klass.__bases__: 

484 return klass 

485 

486 return None 

487 

488 @classmethod 

489 def __get_pydantic_core_schema__(cls, source_type: Any, handler: Any): 

490 if source_type.__name__ == "RegistryBaseModel": 

491 # This class should never be validated directly, but pydantic still builds its schema. 

492 return handler(Any) 

493 elif source_type is cls._find_registered_base(): 

494 # Our direct subclasses, the "user" base model, use the registry's custom discriminator 

495 # schema to validate. 

496 return handler.generate_schema(Annotated[Any, cls.model_registry().discriminator()]) 

497 else: 

498 # Each descendant subclass, members of the registry, are validated as normal. 

499 return handler(source_type)