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

130 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-05-03 02:48 -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 "make_column_literal", 

33) 

34 

35import uuid 

36import warnings 

37from base64 import b64decode, b64encode 

38from functools import cached_property 

39from typing import Literal, TypeAlias, Union, final 

40 

41import astropy.time 

42import erfa 

43from lsst.sphgeom import Region 

44 

45from ..._timespan import Timespan 

46from ...time_utils import TimeConverter 

47from ._base import ColumnLiteralBase 

48 

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

50 

51 

52@final 

53class IntColumnLiteral(ColumnLiteralBase): 

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

55 

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

57 

58 value: int 

59 """The wrapped value.""" 

60 

61 @classmethod 

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

63 """Construct from the wrapped value. 

64 

65 Parameters 

66 ---------- 

67 value : `int` 

68 Value to wrap. 

69 

70 Returns 

71 ------- 

72 expression : `IntColumnLiteral` 

73 Literal expression object. 

74 """ 

75 return cls.model_construct(value=value) 

76 

77 def __str__(self) -> str: 

78 return repr(self.value) 

79 

80 

81@final 

82class StringColumnLiteral(ColumnLiteralBase): 

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

84 

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

86 

87 value: str 

88 """The wrapped value.""" 

89 

90 @classmethod 

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

92 """Construct from the wrapped value. 

93 

94 Parameters 

95 ---------- 

96 value : `str` 

97 Value to wrap. 

98 

99 Returns 

100 ------- 

101 expression : `StrColumnLiteral` 

102 Literal expression object. 

103 """ 

104 return cls.model_construct(value=value) 

105 

106 def __str__(self) -> str: 

107 return repr(self.value) 

108 

109 

110@final 

111class FloatColumnLiteral(ColumnLiteralBase): 

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

113 

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

115 

116 value: float 

117 """The wrapped value.""" 

118 

119 @classmethod 

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

121 """Construct from the wrapped value. 

122 

123 Parameters 

124 ---------- 

125 value : `float` 

126 Value to wrap. 

127 

128 Returns 

129 ------- 

130 expression : `FloatColumnLiteral` 

131 Literal expression object. 

132 """ 

133 return cls.model_construct(value=value) 

134 

135 def __str__(self) -> str: 

136 return repr(self.value) 

137 

138 

139@final 

140class HashColumnLiteral(ColumnLiteralBase): 

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

142 

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

144 use. 

145 """ 

146 

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

148 

149 encoded: bytes 

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

151 

152 @cached_property 

153 def value(self) -> bytes: 

154 """The wrapped value.""" 

155 return b64decode(self.encoded) 

156 

157 @classmethod 

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

159 """Construct from the wrapped value. 

160 

161 Parameters 

162 ---------- 

163 value : `bytes` 

164 Value to wrap. 

165 

166 Returns 

167 ------- 

168 expression : `HashColumnLiteral` 

169 Literal expression object. 

170 """ 

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

172 

173 def __str__(self) -> str: 

174 return "(bytes)" 

175 

176 

177@final 

178class UUIDColumnLiteral(ColumnLiteralBase): 

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

180 

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

182 

183 value: uuid.UUID 

184 """The wrapped value.""" 

185 

186 @classmethod 

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

188 """Construct from the wrapped value. 

189 

190 Parameters 

191 ---------- 

192 value : `uuid.UUID` 

193 Value to wrap. 

194 

195 Returns 

196 ------- 

197 expression : `UUIDColumnLiteral` 

198 Literal expression object. 

199 """ 

200 return cls.model_construct(value=value) 

201 

202 def __str__(self) -> str: 

203 return str(self.value) 

204 

205 

206@final 

207class DateTimeColumnLiteral(ColumnLiteralBase): 

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

209 

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

211 and restored from that on first use. 

212 """ 

213 

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

215 

216 nsec: int 

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

218 

219 @cached_property 

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

221 """The wrapped value.""" 

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

223 

224 @classmethod 

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

226 """Construct from the wrapped value. 

227 

228 Parameters 

229 ---------- 

230 value : `astropy.time.Time` 

231 Value to wrap. 

232 

233 Returns 

234 ------- 

235 expression : `DateTimeColumnLiteral` 

236 Literal expression object. 

237 """ 

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

239 

240 def __str__(self) -> str: 

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

242 # simulated data in the future 

243 with warnings.catch_warnings(): 

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

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

246 

247 

248@final 

249class TimespanColumnLiteral(ColumnLiteralBase): 

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

251 

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

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

254 """ 

255 

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

257 

258 begin_nsec: int 

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

260 (inclusive). 

261 """ 

262 

263 end_nsec: int 

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

265 (exclusive). 

266 """ 

267 

268 @cached_property 

269 def value(self) -> Timespan: 

270 """The wrapped value.""" 

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

272 

273 @classmethod 

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

275 """Construct from the wrapped value. 

276 

277 Parameters 

278 ---------- 

279 value : `..Timespan` 

280 Value to wrap. 

281 

282 Returns 

283 ------- 

284 expression : `TimespanColumnLiteral` 

285 Literal expression object. 

286 """ 

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

288 

289 def __str__(self) -> str: 

290 return str(self.value) 

291 

292 

293@final 

294class RegionColumnLiteral(ColumnLiteralBase): 

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

296 

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

298 first use. 

299 """ 

300 

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

302 

303 encoded: bytes 

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

305 

306 @cached_property 

307 def value(self) -> bytes: 

308 """The wrapped value.""" 

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

310 

311 @classmethod 

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

313 """Construct from the wrapped value. 

314 

315 Parameters 

316 ---------- 

317 value : `..Region` 

318 Value to wrap. 

319 

320 Returns 

321 ------- 

322 expression : `RegionColumnLiteral` 

323 Literal expression object. 

324 """ 

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

326 

327 def __str__(self) -> str: 

328 return "(region)" 

329 

330 

331ColumnLiteral: TypeAlias = Union[ 

332 IntColumnLiteral, 

333 StringColumnLiteral, 

334 FloatColumnLiteral, 

335 HashColumnLiteral, 

336 UUIDColumnLiteral, 

337 DateTimeColumnLiteral, 

338 TimespanColumnLiteral, 

339 RegionColumnLiteral, 

340] 

341 

342 

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

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

345 

346 Parameters 

347 ---------- 

348 value : `LiteralValue` 

349 Value to wrap. 

350 

351 Returns 

352 ------- 

353 expression : `ColumnLiteral` 

354 Literal expression object. 

355 """ 

356 match value: 

357 case int(): 

358 return IntColumnLiteral.from_value(value) 

359 case str(): 

360 return StringColumnLiteral.from_value(value) 

361 case float(): 

362 return FloatColumnLiteral.from_value(value) 

363 case uuid.UUID(): 

364 return UUIDColumnLiteral.from_value(value) 

365 case bytes(): 

366 return HashColumnLiteral.from_value(value) 

367 case astropy.time.Time(): 

368 return DateTimeColumnLiteral.from_value(value) 

369 case Timespan(): 

370 return TimespanColumnLiteral.from_value(value) 

371 case Region(): 

372 return RegionColumnLiteral.from_value(value) 

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