Coverage for python/lsst/daf/butler/queries/tree/_column_literal.py: 69%

130 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 02:51 -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__ = ( 

31 "ColumnLiteral", 

32 "LiteralValue", 

33 "make_column_literal", 

34) 

35 

36import uuid 

37import warnings 

38from base64 import b64decode, b64encode 

39from functools import cached_property 

40from typing import Literal, TypeAlias, Union, final 

41 

42import astropy.time 

43import erfa 

44from lsst.sphgeom import Region 

45 

46from ..._timespan import Timespan 

47from ...time_utils import TimeConverter 

48from ._base import ColumnLiteralBase 

49 

50LiteralValue: TypeAlias = int | str | float | bytes | uuid.UUID | astropy.time.Time | Timespan | Region 

51 

52 

53@final 

54class IntColumnLiteral(ColumnLiteralBase): 

55 """A literal `int` value in a column expression.""" 

56 

57 expression_type: Literal["int"] = "int" 

58 

59 value: int 

60 """The wrapped value.""" 

61 

62 @classmethod 

63 def from_value(cls, value: int) -> IntColumnLiteral: 

64 """Construct from the wrapped value. 

65 

66 Parameters 

67 ---------- 

68 value : `int` 

69 Value to wrap. 

70 

71 Returns 

72 ------- 

73 expression : `IntColumnLiteral` 

74 Literal expression object. 

75 """ 

76 return cls.model_construct(value=value) 

77 

78 def __str__(self) -> str: 

79 return repr(self.value) 

80 

81 

82@final 

83class StringColumnLiteral(ColumnLiteralBase): 

84 """A literal `str` value in a column expression.""" 

85 

86 expression_type: Literal["string"] = "string" 

87 

88 value: str 

89 """The wrapped value.""" 

90 

91 @classmethod 

92 def from_value(cls, value: str) -> StringColumnLiteral: 

93 """Construct from the wrapped value. 

94 

95 Parameters 

96 ---------- 

97 value : `str` 

98 Value to wrap. 

99 

100 Returns 

101 ------- 

102 expression : `StrColumnLiteral` 

103 Literal expression object. 

104 """ 

105 return cls.model_construct(value=value) 

106 

107 def __str__(self) -> str: 

108 return repr(self.value) 

109 

110 

111@final 

112class FloatColumnLiteral(ColumnLiteralBase): 

113 """A literal `float` value in a column expression.""" 

114 

115 expression_type: Literal["float"] = "float" 

116 

117 value: float 

118 """The wrapped value.""" 

119 

120 @classmethod 

121 def from_value(cls, value: float) -> FloatColumnLiteral: 

122 """Construct from the wrapped value. 

123 

124 Parameters 

125 ---------- 

126 value : `float` 

127 Value to wrap. 

128 

129 Returns 

130 ------- 

131 expression : `FloatColumnLiteral` 

132 Literal expression object. 

133 """ 

134 return cls.model_construct(value=value) 

135 

136 def __str__(self) -> str: 

137 return repr(self.value) 

138 

139 

140@final 

141class HashColumnLiteral(ColumnLiteralBase): 

142 """A literal `bytes` value representing a hash in a column expression. 

143 

144 The original value is base64-encoded when serialized and decoded on first 

145 use. 

146 """ 

147 

148 expression_type: Literal["hash"] = "hash" 

149 

150 encoded: bytes 

151 """The wrapped value after base64 encoding.""" 

152 

153 @cached_property 

154 def value(self) -> bytes: 

155 """The wrapped value.""" 

156 return b64decode(self.encoded) 

157 

158 @classmethod 

159 def from_value(cls, value: bytes) -> HashColumnLiteral: 

160 """Construct from the wrapped value. 

161 

162 Parameters 

163 ---------- 

164 value : `bytes` 

165 Value to wrap. 

166 

167 Returns 

168 ------- 

169 expression : `HashColumnLiteral` 

170 Literal expression object. 

171 """ 

172 return cls.model_construct(encoded=b64encode(value)) 

173 

174 def __str__(self) -> str: 

175 return "(bytes)" 

176 

177 

178@final 

179class UUIDColumnLiteral(ColumnLiteralBase): 

180 """A literal `uuid.UUID` value in a column expression.""" 

181 

182 expression_type: Literal["uuid"] = "uuid" 

183 

184 value: uuid.UUID 

185 """The wrapped value.""" 

186 

187 @classmethod 

188 def from_value(cls, value: uuid.UUID) -> UUIDColumnLiteral: 

189 """Construct from the wrapped value. 

190 

191 Parameters 

192 ---------- 

193 value : `uuid.UUID` 

194 Value to wrap. 

195 

196 Returns 

197 ------- 

198 expression : `UUIDColumnLiteral` 

199 Literal expression object. 

200 """ 

201 return cls.model_construct(value=value) 

202 

203 def __str__(self) -> str: 

204 return str(self.value) 

205 

206 

207@final 

208class DateTimeColumnLiteral(ColumnLiteralBase): 

209 """A literal `astropy.time.Time` value in a column expression. 

210 

211 The time is converted into TAI nanoseconds since 1970-01-01 when serialized 

212 and restored from that on first use. 

213 """ 

214 

215 expression_type: Literal["datetime"] = "datetime" 

216 

217 nsec: int 

218 """TAI nanoseconds since 1970-01-01.""" 

219 

220 @cached_property 

221 def value(self) -> astropy.time.Time: 

222 """The wrapped value.""" 

223 return TimeConverter().nsec_to_astropy(self.nsec) 

224 

225 @classmethod 

226 def from_value(cls, value: astropy.time.Time) -> DateTimeColumnLiteral: 

227 """Construct from the wrapped value. 

228 

229 Parameters 

230 ---------- 

231 value : `astropy.time.Time` 

232 Value to wrap. 

233 

234 Returns 

235 ------- 

236 expression : `DateTimeColumnLiteral` 

237 Literal expression object. 

238 """ 

239 return cls.model_construct(nsec=TimeConverter().astropy_to_nsec(value)) 

240 

241 def __str__(self) -> str: 

242 # Trap dubious year warnings in case we have timespans from 

243 # simulated data in the future 

244 with warnings.catch_warnings(): 

245 warnings.simplefilter("ignore", category=erfa.ErfaWarning) 

246 return self.value.tai.strftime("%Y-%m-%dT%H:%M:%S") 

247 

248 

249@final 

250class TimespanColumnLiteral(ColumnLiteralBase): 

251 """A literal `Timespan` value in a column expression. 

252 

253 The timespan bounds are converted into TAI nanoseconds since 1970-01-01 

254 when serialized and the timespan is restored from that on first use. 

255 """ 

256 

257 expression_type: Literal["timespan"] = "timespan" 

258 

259 begin_nsec: int 

260 """TAI nanoseconds since 1970-01-01 for the lower bound of the timespan 

261 (inclusive). 

262 """ 

263 

264 end_nsec: int 

265 """TAI nanoseconds since 1970-01-01 for the upper bound of the timespan 

266 (exclusive). 

267 """ 

268 

269 @cached_property 

270 def value(self) -> Timespan: 

271 """The wrapped value.""" 

272 return Timespan(None, None, _nsec=(self.begin_nsec, self.end_nsec)) 

273 

274 @classmethod 

275 def from_value(cls, value: Timespan) -> TimespanColumnLiteral: 

276 """Construct from the wrapped value. 

277 

278 Parameters 

279 ---------- 

280 value : `..Timespan` 

281 Value to wrap. 

282 

283 Returns 

284 ------- 

285 expression : `TimespanColumnLiteral` 

286 Literal expression object. 

287 """ 

288 return cls.model_construct(begin_nsec=value.nsec[0], end_nsec=value.nsec[1]) 

289 

290 def __str__(self) -> str: 

291 return str(self.value) 

292 

293 

294@final 

295class RegionColumnLiteral(ColumnLiteralBase): 

296 """A literal `lsst.sphgeom.Region` value in a column expression. 

297 

298 The region is encoded to base64 `bytes` when serialized, and decoded on 

299 first use. 

300 """ 

301 

302 expression_type: Literal["region"] = "region" 

303 

304 encoded: bytes 

305 """The wrapped value after base64 encoding.""" 

306 

307 @cached_property 

308 def value(self) -> bytes: 

309 """The wrapped value.""" 

310 return Region.decode(b64decode(self.encoded)) 

311 

312 @classmethod 

313 def from_value(cls, value: Region) -> RegionColumnLiteral: 

314 """Construct from the wrapped value. 

315 

316 Parameters 

317 ---------- 

318 value : `..Region` 

319 Value to wrap. 

320 

321 Returns 

322 ------- 

323 expression : `RegionColumnLiteral` 

324 Literal expression object. 

325 """ 

326 return cls.model_construct(encoded=b64encode(value.encode())) 

327 

328 def __str__(self) -> str: 

329 return "(region)" 

330 

331 

332ColumnLiteral: TypeAlias = Union[ 

333 IntColumnLiteral, 

334 StringColumnLiteral, 

335 FloatColumnLiteral, 

336 HashColumnLiteral, 

337 UUIDColumnLiteral, 

338 DateTimeColumnLiteral, 

339 TimespanColumnLiteral, 

340 RegionColumnLiteral, 

341] 

342 

343 

344def make_column_literal(value: LiteralValue) -> ColumnLiteral: 

345 """Construct a `ColumnLiteral` from the value it will wrap. 

346 

347 Parameters 

348 ---------- 

349 value : `LiteralValue` 

350 Value to wrap. 

351 

352 Returns 

353 ------- 

354 expression : `ColumnLiteral` 

355 Literal expression object. 

356 """ 

357 match value: 

358 case int(): 

359 return IntColumnLiteral.from_value(value) 

360 case str(): 

361 return StringColumnLiteral.from_value(value) 

362 case float(): 

363 return FloatColumnLiteral.from_value(value) 

364 case uuid.UUID(): 

365 return UUIDColumnLiteral.from_value(value) 

366 case bytes(): 

367 return HashColumnLiteral.from_value(value) 

368 case astropy.time.Time(): 

369 return DateTimeColumnLiteral.from_value(value) 

370 case Timespan(): 

371 return TimespanColumnLiteral.from_value(value) 

372 case Region(): 

373 return RegionColumnLiteral.from_value(value) 

374 raise TypeError(f"Invalid type {type(value).__name__!r} of value {value!r} for column literal.")