Coverage for python/lsst/daf/butler/pydantic_utils.py: 48%

75 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-30 02:53 -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/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ("DeferredValidation", "get_universe_from_context", "SerializableRegion", "SerializableTime") 

31 

32from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Generic, Self, TypeAlias, TypeVar, get_args 

33 

34import pydantic 

35from astropy.time import Time 

36from lsst.sphgeom import Region 

37from pydantic_core import core_schema 

38 

39from .time_utils import TimeConverter 

40 

41if TYPE_CHECKING: 

42 from .dimensions import DimensionUniverse 

43 

44_T = TypeVar("_T") 

45 

46 

47def get_universe_from_context(context: dict[str, Any] | None) -> DimensionUniverse: 

48 """Extract the dimension universe from a Pydantic validation context 

49 dictionary. 

50 

51 Parameters 

52 ---------- 

53 context : `dict` 

54 Dictionary obtained from `pydantic.ValidationInfo.context`. 

55 

56 Returns 

57 ------- 

58 universe : `DimensionUniverse` 

59 Definitions for all dimensions. 

60 

61 Notes 

62 ----- 

63 This function just provides consistent error handling around:: 

64 

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 

76 

77 

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. 

81 

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. 

91 

92 Notes 

93 ----- 

94 This class must be subclassed to be used, but subclasses are always 

95 trivial:: 

96 

97 class SerializableThing(DeferredValidation[Thing]): 

98 pass 

99 

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. 

105 

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`. 

116 

117 Indirect subclasses of `DeferredValidation` are not permitted. 

118 

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 """ 

125 

126 def __init__(self, data: Any): 

127 self._data = data 

128 self._is_validated = False 

129 

130 @classmethod 

131 def from_validated(cls, wrapped: _T) -> Self: 

132 """Construct from an instance of the wrapped type. 

133 

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`. 

139 

140 Parameters 

141 ---------- 

142 wrapped : `object` 

143 Instance of the wrapped type. 

144 

145 Returns 

146 ------- 

147 wrapper : `DeferredValidation` 

148 Instance of the wrapper. 

149 """ 

150 result = cls(wrapped) 

151 result._is_validated = True 

152 return result 

153 

154 def validated(self, **kwargs: Any) -> _T: 

155 """Validate (if necessary) and return the validated object. 

156 

157 Parameters 

158 ---------- 

159 **kwargs 

160 Additional keywords arguments are passed as the Pydantic 

161 "validation context" `dict`. 

162 

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 

175 

176 _WRAPPED_TYPE: ClassVar[Any | None] = None 

177 _WRAPPED_TYPE_ADAPTER: ClassVar[pydantic.TypeAdapter[Any] | None] = None 

178 

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__() 

197 

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 

208 

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 

215 

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 ) 

242 

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) 

252 

253 

254def _deserialize_region(value: object, handler: pydantic.ValidatorFunctionWrapHandler) -> Region: 

255 if isinstance(value, Region): 

256 return value 

257 

258 string = handler(value) 

259 return Region.decode(bytes.fromhex(string)) 

260 

261 

262def _serialize_region(region: Region) -> str: 

263 return region.encode().hex() 

264 

265 

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`. 

280 

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""" 

286 

287 

288def _deserialize_time(value: object, handler: pydantic.ValidatorFunctionWrapHandler) -> Region: 

289 if isinstance(value, Time): 

290 return value 

291 

292 integer = handler(value) 

293 return TimeConverter().nsec_to_astropy(integer) 

294 

295 

296def _serialize_time(time: Time) -> int: 

297 return TimeConverter().astropy_to_nsec(time) 

298 

299 

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`. 

313 

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"""