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

75 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-13 09:58 +0000

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(self._data, context=kwargs) 

171 self._is_validated = True 

172 return self._data 

173 

174 _WRAPPED_TYPE: ClassVar[Any | None] = None 

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

176 

177 def __init_subclass__(cls) -> None: 

178 # We override __init_subclass__ to grab the type argument to the 

179 # DeferredValidation base class, since that's the wrapped type. 

180 assert ( 

181 cls.__base__ is DeferredValidation 

182 ), "Indirect subclasses of DeferredValidation are not allowed." 

183 try: 

184 # This uses some typing internals that are not as stable as the 

185 # rest of Python, so it's the messiest aspect of this class, but 

186 # even if it breaks on (say) some Python minor releases, it should 

187 # be easy to detect and fix and I think that makes it better than 

188 # requiring the wrapped type to be declared twice when subclassing. 

189 # Since the type-checking ecosystem depends on this sort of thing 

190 # to work it's not exactly private, either. 

191 cls._WRAPPED_TYPE = get_args(cls.__orig_bases__[0])[0] # type: ignore 

192 except Exception as err: 

193 raise TypeError("DeferredValidation must be subclassed with a single type parameter.") from err 

194 return super().__init_subclass__() 

195 

196 @classmethod 

197 def _get_wrapped_type_adapter(cls) -> pydantic.TypeAdapter[_T]: 

198 """Return the Pydantic adapter for the wrapped type, constructing and 

199 caching it if necessary. 

200 """ 

201 if cls._WRAPPED_TYPE_ADAPTER is None: 

202 if cls._WRAPPED_TYPE is None: 

203 raise TypeError("DeferredValidation must be subclassed to be used.") 

204 cls._WRAPPED_TYPE_ADAPTER = pydantic.TypeAdapter(cls._WRAPPED_TYPE) 

205 return cls._WRAPPED_TYPE_ADAPTER 

206 

207 def _serialize(self) -> Any: 

208 """Serialize this object.""" 

209 if self._is_validated: 

210 return self._get_wrapped_type_adapter().dump_python(self._data) 

211 else: 

212 return self._data 

213 

214 @classmethod 

215 def __get_pydantic_core_schema__( 

216 cls, _source_type: Any, _handler: pydantic.GetCoreSchemaHandler 

217 ) -> core_schema.CoreSchema: 

218 # This is the Pydantic hook for overriding serialization and 

219 # validation. It's also normally the hook for defining the JSON 

220 # schema, but we throw that JSON schema away and define our own in 

221 # __get_pydantic_json_schema__. 

222 return core_schema.json_or_python_schema( 

223 # When deserializing from JSON, invoke the constructor with the 

224 # result of parsing the JSON into Python primitives. 

225 json_schema=core_schema.no_info_plain_validator_function(cls), 

226 # When validating a Python dict... 

227 python_schema=core_schema.union_schema( 

228 [ 

229 # ...first see if we already have an instance of the 

230 # wrapper... 

231 core_schema.is_instance_schema(cls), 

232 # ...and otherwise just call the constructor on whatever 

233 # we were given. 

234 core_schema.no_info_plain_validator_function(cls), 

235 ] 

236 ), 

237 # When serializing to JSON, just call the _serialize method. 

238 serialization=core_schema.plain_serializer_function_ser_schema(cls._serialize), 

239 ) 

240 

241 @classmethod 

242 def __get_pydantic_json_schema__( 

243 cls, _core_schema: core_schema.CoreSchema, handler: pydantic.json_schema.GetJsonSchemaHandler 

244 ) -> pydantic.json_schema.JsonSchemaValue: 

245 # This is the Pydantic hook for customizing JSON schema. We ignore 

246 # the schema generated for this class, and just return the JSON schema 

247 # of the wrapped type. 

248 json_schema = handler(cls._get_wrapped_type_adapter().core_schema) 

249 return handler.resolve_ref_schema(json_schema) 

250 

251 

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

253 if isinstance(value, Region): 

254 return value 

255 

256 string = handler(value) 

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

258 

259 

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

261 return region.encode().hex() 

262 

263 

264SerializableRegion: TypeAlias = Annotated[ 264 ↛ exitline 264 didn't jump to the function exit

265 Region, 

266 pydantic.GetPydanticSchema(lambda _, h: h(str)), 

267 pydantic.WrapValidator(_deserialize_region), 

268 pydantic.PlainSerializer(_serialize_region), 

269 pydantic.WithJsonSchema( 

270 { 

271 "type": "string", 

272 "description": "A region on the sphere from the lsst.sphgeom package.", 

273 "media": {"binaryEncoding": "base16", "type": "application/lsst.sphgeom"}, 

274 } 

275 ), 

276] 

277"""A Pydantic-annotated version of `lsst.sphgeom.Region`. 

278 

279An object annotated with this type is always an `lsst.sphgeom.Region` instance 

280in Python, but unlike `lsst.sphgeom.Region` itself it can be used as a type 

281in Pydantic models and type adapters, resulting in the field being saved as 

282a hex encoding of the sphgeom-encoded bytes. 

283""" 

284 

285 

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

287 if isinstance(value, Time): 

288 return value 

289 

290 integer = handler(value) 

291 return TimeConverter().nsec_to_astropy(integer) 

292 

293 

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

295 return TimeConverter().astropy_to_nsec(time) 

296 

297 

298SerializableTime: TypeAlias = Annotated[ 298 ↛ exitline 298 didn't jump to the function exit

299 Time, 

300 pydantic.GetPydanticSchema(lambda _, h: h(int)), 

301 pydantic.WrapValidator(_deserialize_time), 

302 pydantic.PlainSerializer(_serialize_time), 

303 pydantic.WithJsonSchema( 

304 { 

305 "type": "integer", 

306 "description": "A TAI time represented as integer nanoseconds since 1970-01-01 00:00:00.", 

307 } 

308 ), 

309] 

310"""A Pydantic-annotated version of `astropy.time.Time`. 

311 

312An object annotated with this type is always an `astropy.time.Time` instance 

313in Python, but unlike `astropy.time.Time` itself it can be used as a type 

314in Pydantic models and type adapters, resulting in the field being saved as 

315integer nanoseconds since 1970-01-01 00:00:00. 

316"""