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

54 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-10 10:14 +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") 

31 

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

33 

34import pydantic 

35from pydantic_core import core_schema 

36 

37if TYPE_CHECKING: 

38 from .dimensions import DimensionUniverse 

39 

40_T = TypeVar("_T") 

41 

42 

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

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

45 dictionary. 

46 

47 Parameters 

48 ---------- 

49 context : `dict` 

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

51 

52 Returns 

53 ------- 

54 universe : `DimensionUniverse` 

55 Definitions for all dimensions. 

56 

57 Notes 

58 ----- 

59 This function just provides consistent error handling around:: 

60 

61 context["universe"] 

62 """ 

63 if context is None: 

64 raise ValueError("This object requires Pydantic validation context to be deserialized.") 

65 try: 

66 return context["universe"] 

67 except KeyError: 

68 raise ValueError( 

69 "This object requires the DimensionUniverse to be provided in the Pydantic validation " 

70 "context to be deserialized." 

71 ) from None 

72 

73 

74class DeferredValidation(Generic[_T]): 

75 """A base class whose subclasses define a wrapper for a Pydantic-aware type 

76 that defers validation but declares the same JSON schema. 

77 

78 Parameters 

79 ---------- 

80 data : `object` 

81 Unvalidated data representing an instance of the wrapped type. This 

82 may be the serialized form of the wrapped type, an instance of the 

83 wrapped type, or anything else - but the in the latter case, calls to 

84 `validated` will fail with a Pydantic validation error, and if the 

85 object is known to be an instance of the wrapped type, `from_validated` 

86 should be preferred. 

87 

88 Notes 

89 ----- 

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

91 trivial:: 

92 

93 class SerializableThing(DeferredValidation[Thing]): 

94 pass 

95 

96 The type parameter for `DeferredValidation` may be a special typing object 

97 like `typing.Union` or `typing.Annotated` instead of an actual `type` 

98 object. The only requirement is that it must be a type Pydantic 

99 recognizes, like a `pydantic.BaseModel` subclass, a dataclass, or a 

100 primitive built-in. 

101 

102 A wrapper subclass (e.g. ``SerializableThing``) can be used with Pydantic 

103 via `pydantic.TypeAdapter` or as a field in `pydantic.BaseModel`. The JSON 

104 schema of the wrapper will be consistent with the JSON schema of the 

105 wrapped type (though it may not use JSON pointer references the same way), 

106 and Pydantic serialization will work regardless of whether the wrapper 

107 instance was initialized with the raw type or the wrapped type. Pydantic 

108 validation of the wrapper will effectively do nothing, however; instead, 

109 the `validated` method must be called to return a fully-validated instance 

110 of the wrapped type, which is then cached within the wrapper for subsequent 

111 calls to `validated`. 

112 

113 Indirect subclasses of `DeferredValidation` are not permitted. 

114 

115 A major use case for `DeferredValidation` is types whose validation 

116 requires additional runtime context (via the Pydantic "validation context" 

117 dictionary that can custom validator hooks can access). These types are 

118 often first deserialized (e.g. by FastAPI) in a way that does not permit 

119 that context to be provided. 

120 """ 

121 

122 def __init__(self, data: Any): 

123 self._data = data 

124 self._is_validated = False 

125 

126 @classmethod 

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

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

129 

130 Unlike invoking the constructor with an instance of the wrapped type, 

131 this factory marks the held instance as already validated (since that 

132 is expected to be guaranteed by the caller, possibly with the help of 

133 static analysis), which sidesteps Pydantic validation in later calls 

134 to `validated`. 

135 

136 Parameters 

137 ---------- 

138 wrapped : `object` 

139 Instance of the wrapped type. 

140 

141 Returns 

142 ------- 

143 wrapper : `DeferredValidation` 

144 Instance of the wrapper. 

145 """ 

146 result = cls(wrapped) 

147 result._is_validated = True 

148 return result 

149 

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

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

152 

153 Parameters 

154 ---------- 

155 **kwargs 

156 Additional keywords arguments are passed as the Pydantic 

157 "validation context" `dict`. 

158 

159 Returns 

160 ------- 

161 wrapped 

162 An instance of the wrapped type. This is also cached for the next 

163 call to `validated`, *which will ignore ``**kwargs``*. 

164 """ 

165 if not self._is_validated: 

166 self._data = self._get_wrapped_type_adapter().validate_python(self._data, context=kwargs) 

167 self._is_validated = True 

168 return self._data 

169 

170 _WRAPPED_TYPE: ClassVar[Any | None] = None 

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

172 

173 def __init_subclass__(cls) -> None: 

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

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

176 assert ( 

177 cls.__base__ is DeferredValidation 

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

179 try: 

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

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

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

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

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

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

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

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

188 except Exception as err: 

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

190 return super().__init_subclass__() 

191 

192 @classmethod 

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

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

195 caching it if necessary. 

196 """ 

197 if cls._WRAPPED_TYPE_ADAPTER is None: 

198 if cls._WRAPPED_TYPE is None: 

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

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

201 return cls._WRAPPED_TYPE_ADAPTER 

202 

203 def _serialize(self) -> Any: 

204 """Serialize this object.""" 

205 if self._is_validated: 

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

207 else: 

208 return self._data 

209 

210 @classmethod 

211 def __get_pydantic_core_schema__( 

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

213 ) -> core_schema.CoreSchema: 

214 # This is the Pydantic hook for overriding serialization and 

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

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

217 # __get_pydantic_json_schema__. 

218 return core_schema.json_or_python_schema( 

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

220 # result of parsing the JSON into Python primitives. 

221 json_schema=core_schema.no_info_plain_validator_function(cls), 

222 # When validating a Python dict... 

223 python_schema=core_schema.union_schema( 

224 [ 

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

226 # wrapper... 

227 core_schema.is_instance_schema(cls), 

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

229 # we were given. 

230 core_schema.no_info_plain_validator_function(cls), 

231 ] 

232 ), 

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

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

235 ) 

236 

237 @classmethod 

238 def __get_pydantic_json_schema__( 

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

240 ) -> pydantic.json_schema.JsonSchemaValue: 

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

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

243 # of the wrapped type. 

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

245 return handler.resolve_ref_schema(json_schema)