Coverage for python/lsst/daf/butler/pydantic_utils.py: 48%
75 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-07 02:46 -0700
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-07 02:46 -0700
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = ("DeferredValidation", "get_universe_from_context", "SerializableRegion", "SerializableTime")
32from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Generic, Self, TypeAlias, TypeVar, get_args
34import pydantic
35from astropy.time import Time
36from lsst.sphgeom import Region
37from pydantic_core import core_schema
39from .time_utils import TimeConverter
41if TYPE_CHECKING:
42 from .dimensions import DimensionUniverse
44_T = TypeVar("_T")
47def get_universe_from_context(context: dict[str, Any] | None) -> DimensionUniverse:
48 """Extract the dimension universe from a Pydantic validation context
49 dictionary.
51 Parameters
52 ----------
53 context : `dict`
54 Dictionary obtained from `pydantic.ValidationInfo.context`.
56 Returns
57 -------
58 universe : `DimensionUniverse`
59 Definitions for all dimensions.
61 Notes
62 -----
63 This function just provides consistent error handling around::
65 context["universe"]
66 """
67 if context is None:
68 raise ValueError("This object requires Pydantic validation context to be deserialized.")
69 try:
70 return context["universe"]
71 except KeyError:
72 raise ValueError(
73 "This object requires the DimensionUniverse to be provided in the Pydantic validation "
74 "context to be deserialized."
75 ) from None
78class DeferredValidation(Generic[_T]):
79 """A base class whose subclasses define a wrapper for a Pydantic-aware type
80 that defers validation but declares the same JSON schema.
82 Parameters
83 ----------
84 data : `object`
85 Unvalidated data representing an instance of the wrapped type. This
86 may be the serialized form of the wrapped type, an instance of the
87 wrapped type, or anything else - but the in the latter case, calls to
88 `validated` will fail with a Pydantic validation error, and if the
89 object is known to be an instance of the wrapped type, `from_validated`
90 should be preferred.
92 Notes
93 -----
94 This class must be subclassed to be used, but subclasses are always
95 trivial::
97 class SerializableThing(DeferredValidation[Thing]):
98 pass
100 The type parameter for `DeferredValidation` may be a special typing object
101 like `typing.Union` or `typing.Annotated` instead of an actual `type`
102 object. The only requirement is that it must be a type Pydantic
103 recognizes, like a `pydantic.BaseModel` subclass, a dataclass, or a
104 primitive built-in.
106 A wrapper subclass (e.g. ``SerializableThing``) can be used with Pydantic
107 via `pydantic.TypeAdapter` or as a field in `pydantic.BaseModel`. The JSON
108 schema of the wrapper will be consistent with the JSON schema of the
109 wrapped type (though it may not use JSON pointer references the same way),
110 and Pydantic serialization will work regardless of whether the wrapper
111 instance was initialized with the raw type or the wrapped type. Pydantic
112 validation of the wrapper will effectively do nothing, however; instead,
113 the `validated` method must be called to return a fully-validated instance
114 of the wrapped type, which is then cached within the wrapper for subsequent
115 calls to `validated`.
117 Indirect subclasses of `DeferredValidation` are not permitted.
119 A major use case for `DeferredValidation` is types whose validation
120 requires additional runtime context (via the Pydantic "validation context"
121 dictionary that can custom validator hooks can access). These types are
122 often first deserialized (e.g. by FastAPI) in a way that does not permit
123 that context to be provided.
124 """
126 def __init__(self, data: Any):
127 self._data = data
128 self._is_validated = False
130 @classmethod
131 def from_validated(cls, wrapped: _T) -> Self:
132 """Construct from an instance of the wrapped type.
134 Unlike invoking the constructor with an instance of the wrapped type,
135 this factory marks the held instance as already validated (since that
136 is expected to be guaranteed by the caller, possibly with the help of
137 static analysis), which sidesteps Pydantic validation in later calls
138 to `validated`.
140 Parameters
141 ----------
142 wrapped : `object`
143 Instance of the wrapped type.
145 Returns
146 -------
147 wrapper : `DeferredValidation`
148 Instance of the wrapper.
149 """
150 result = cls(wrapped)
151 result._is_validated = True
152 return result
154 def validated(self, **kwargs: Any) -> _T:
155 """Validate (if necessary) and return the validated object.
157 Parameters
158 ----------
159 **kwargs
160 Additional keywords arguments are passed as the Pydantic
161 "validation context" `dict`.
163 Returns
164 -------
165 wrapped
166 An instance of the wrapped type. This is also cached for the next
167 call to `validated`, *which will ignore ``**kwargs``*.
168 """
169 if not self._is_validated:
170 self._data = self._get_wrapped_type_adapter().validate_python(
171 self._data, strict=False, context=kwargs
172 )
173 self._is_validated = True
174 return self._data
176 _WRAPPED_TYPE: ClassVar[Any | None] = None
177 _WRAPPED_TYPE_ADAPTER: ClassVar[pydantic.TypeAdapter[Any] | None] = None
179 def __init_subclass__(cls) -> None:
180 # We override __init_subclass__ to grab the type argument to the
181 # DeferredValidation base class, since that's the wrapped type.
182 assert (
183 cls.__base__ is DeferredValidation
184 ), "Indirect subclasses of DeferredValidation are not allowed."
185 try:
186 # This uses some typing internals that are not as stable as the
187 # rest of Python, so it's the messiest aspect of this class, but
188 # even if it breaks on (say) some Python minor releases, it should
189 # be easy to detect and fix and I think that makes it better than
190 # requiring the wrapped type to be declared twice when subclassing.
191 # Since the type-checking ecosystem depends on this sort of thing
192 # to work it's not exactly private, either.
193 cls._WRAPPED_TYPE = get_args(cls.__orig_bases__[0])[0] # type: ignore
194 except Exception as err:
195 raise TypeError("DeferredValidation must be subclassed with a single type parameter.") from err
196 return super().__init_subclass__()
198 @classmethod
199 def _get_wrapped_type_adapter(cls) -> pydantic.TypeAdapter[_T]:
200 """Return the Pydantic adapter for the wrapped type, constructing and
201 caching it if necessary.
202 """
203 if cls._WRAPPED_TYPE_ADAPTER is None:
204 if cls._WRAPPED_TYPE is None:
205 raise TypeError("DeferredValidation must be subclassed to be used.")
206 cls._WRAPPED_TYPE_ADAPTER = pydantic.TypeAdapter(cls._WRAPPED_TYPE)
207 return cls._WRAPPED_TYPE_ADAPTER
209 def _serialize(self) -> Any:
210 """Serialize this object."""
211 if self._is_validated:
212 return self._get_wrapped_type_adapter().dump_python(self._data)
213 else:
214 return self._data
216 @classmethod
217 def __get_pydantic_core_schema__(
218 cls, _source_type: Any, _handler: pydantic.GetCoreSchemaHandler
219 ) -> core_schema.CoreSchema:
220 # This is the Pydantic hook for overriding serialization and
221 # validation. It's also normally the hook for defining the JSON
222 # schema, but we throw that JSON schema away and define our own in
223 # __get_pydantic_json_schema__.
224 return core_schema.json_or_python_schema(
225 # When deserializing from JSON, invoke the constructor with the
226 # result of parsing the JSON into Python primitives.
227 json_schema=core_schema.no_info_plain_validator_function(cls),
228 # When validating a Python dict...
229 python_schema=core_schema.union_schema(
230 [
231 # ...first see if we already have an instance of the
232 # wrapper...
233 core_schema.is_instance_schema(cls),
234 # ...and otherwise just call the constructor on whatever
235 # we were given.
236 core_schema.no_info_plain_validator_function(cls),
237 ]
238 ),
239 # When serializing to JSON, just call the _serialize method.
240 serialization=core_schema.plain_serializer_function_ser_schema(cls._serialize),
241 )
243 @classmethod
244 def __get_pydantic_json_schema__(
245 cls, _core_schema: core_schema.CoreSchema, handler: pydantic.json_schema.GetJsonSchemaHandler
246 ) -> pydantic.json_schema.JsonSchemaValue:
247 # This is the Pydantic hook for customizing JSON schema. We ignore
248 # the schema generated for this class, and just return the JSON schema
249 # of the wrapped type.
250 json_schema = handler(cls._get_wrapped_type_adapter().core_schema)
251 return handler.resolve_ref_schema(json_schema)
254def _deserialize_region(value: object, handler: pydantic.ValidatorFunctionWrapHandler) -> Region:
255 if isinstance(value, Region):
256 return value
258 string = handler(value)
259 return Region.decode(bytes.fromhex(string))
262def _serialize_region(region: Region) -> str:
263 return region.encode().hex()
266SerializableRegion: TypeAlias = Annotated[ 266 ↛ exitline 266 didn't jump to the function exit
267 Region,
268 pydantic.GetPydanticSchema(lambda _, h: h(str)),
269 pydantic.WrapValidator(_deserialize_region),
270 pydantic.PlainSerializer(_serialize_region),
271 pydantic.WithJsonSchema(
272 {
273 "type": "string",
274 "description": "A region on the sphere from the lsst.sphgeom package.",
275 "media": {"binaryEncoding": "base16", "type": "application/lsst.sphgeom"},
276 }
277 ),
278]
279"""A Pydantic-annotated version of `lsst.sphgeom.Region`.
281An object annotated with this type is always an `lsst.sphgeom.Region` instance
282in Python, but unlike `lsst.sphgeom.Region` itself it can be used as a type
283in Pydantic models and type adapters, resulting in the field being saved as
284a hex encoding of the sphgeom-encoded bytes.
285"""
288def _deserialize_time(value: object, handler: pydantic.ValidatorFunctionWrapHandler) -> Region:
289 if isinstance(value, Time):
290 return value
292 integer = handler(value)
293 return TimeConverter().nsec_to_astropy(integer)
296def _serialize_time(time: Time) -> int:
297 return TimeConverter().astropy_to_nsec(time)
300SerializableTime: TypeAlias = Annotated[ 300 ↛ exitline 300 didn't jump to the function exit
301 Time,
302 pydantic.GetPydanticSchema(lambda _, h: h(int)),
303 pydantic.WrapValidator(_deserialize_time),
304 pydantic.PlainSerializer(_serialize_time),
305 pydantic.WithJsonSchema(
306 {
307 "type": "integer",
308 "description": "A TAI time represented as integer nanoseconds since 1970-01-01 00:00:00.",
309 }
310 ),
311]
312"""A Pydantic-annotated version of `astropy.time.Time`.
314An object annotated with this type is always an `astropy.time.Time` instance
315in Python, but unlike `astropy.time.Time` itself it can be used as a type
316in Pydantic models and type adapters, resulting in the field being saved as
317integer nanoseconds since 1970-01-01 00:00:00.
318"""