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

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 datetime 

37import numbers 

38import uuid 

39import warnings 

40from base64 import b64decode, b64encode 

41from functools import cached_property 

42from typing import Literal, TypeAlias, Union, final 

43 

44import astropy.coordinates 

45import astropy.time 

46import erfa 

47 

48import lsst.sphgeom 

49 

50from ..._timespan import Timespan 

51from ...time_utils import TimeConverter 

52from ._base import ColumnLiteralBase 

53 

54LiteralValue: TypeAlias = ( 

55 int 

56 | str 

57 | float 

58 | bytes 

59 | uuid.UUID 

60 | astropy.time.Time 

61 | datetime.datetime 

62 | Timespan 

63 | lsst.sphgeom.Region 

64 | lsst.sphgeom.LonLat 

65) 

66 

67 

68@final 

69class IntColumnLiteral(ColumnLiteralBase): 

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

71 

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

73 

74 value: int 

75 """The wrapped value.""" 

76 

77 @classmethod 

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

79 """Construct from the wrapped value. 

80 

81 Parameters 

82 ---------- 

83 value : `int` 

84 Value to wrap. 

85 

86 Returns 

87 ------- 

88 expression : `IntColumnLiteral` 

89 Literal expression object. 

90 """ 

91 return cls.model_construct(value=value) 

92 

93 def __str__(self) -> str: 

94 return repr(self.value) 

95 

96 

97@final 

98class StringColumnLiteral(ColumnLiteralBase): 

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

100 

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

102 

103 value: str 

104 """The wrapped value.""" 

105 

106 @classmethod 

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

108 """Construct from the wrapped value. 

109 

110 Parameters 

111 ---------- 

112 value : `str` 

113 Value to wrap. 

114 

115 Returns 

116 ------- 

117 expression : `StrColumnLiteral` 

118 Literal expression object. 

119 """ 

120 return cls.model_construct(value=value) 

121 

122 def __str__(self) -> str: 

123 return repr(self.value) 

124 

125 

126@final 

127class FloatColumnLiteral(ColumnLiteralBase): 

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

129 

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

131 

132 value: float 

133 """The wrapped value.""" 

134 

135 @classmethod 

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

137 """Construct from the wrapped value. 

138 

139 Parameters 

140 ---------- 

141 value : `float` 

142 Value to wrap. 

143 

144 Returns 

145 ------- 

146 expression : `FloatColumnLiteral` 

147 Literal expression object. 

148 """ 

149 return cls.model_construct(value=value) 

150 

151 def __str__(self) -> str: 

152 return repr(self.value) 

153 

154 

155@final 

156class HashColumnLiteral(ColumnLiteralBase): 

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

158 

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

160 use. 

161 """ 

162 

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

164 

165 encoded: bytes 

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

167 

168 @cached_property 

169 def value(self) -> bytes: 

170 """The wrapped value.""" 

171 return b64decode(self.encoded) 

172 

173 @classmethod 

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

175 """Construct from the wrapped value. 

176 

177 Parameters 

178 ---------- 

179 value : `bytes` 

180 Value to wrap. 

181 

182 Returns 

183 ------- 

184 expression : `HashColumnLiteral` 

185 Literal expression object. 

186 """ 

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

188 

189 def __str__(self) -> str: 

190 return "(bytes)" 

191 

192 

193@final 

194class UUIDColumnLiteral(ColumnLiteralBase): 

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

196 

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

198 

199 value: uuid.UUID 

200 """The wrapped value.""" 

201 

202 @classmethod 

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

204 """Construct from the wrapped value. 

205 

206 Parameters 

207 ---------- 

208 value : `uuid.UUID` 

209 Value to wrap. 

210 

211 Returns 

212 ------- 

213 expression : `UUIDColumnLiteral` 

214 Literal expression object. 

215 """ 

216 return cls.model_construct(value=value) 

217 

218 def __str__(self) -> str: 

219 return str(self.value) 

220 

221 

222@final 

223class DateTimeColumnLiteral(ColumnLiteralBase): 

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

225 

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

227 and restored from that on first use. 

228 """ 

229 

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

231 

232 nsec: int 

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

234 

235 @cached_property 

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

237 """The wrapped value.""" 

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

239 

240 @classmethod 

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

242 """Construct from the wrapped value. 

243 

244 Parameters 

245 ---------- 

246 value : `astropy.time.Time` 

247 Value to wrap. 

248 

249 Returns 

250 ------- 

251 expression : `DateTimeColumnLiteral` 

252 Literal expression object. 

253 """ 

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

255 

256 def __str__(self) -> str: 

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

258 # simulated data in the future 

259 with warnings.catch_warnings(): 

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

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

262 

263 

264@final 

265class TimespanColumnLiteral(ColumnLiteralBase): 

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

267 

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

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

270 """ 

271 

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

273 

274 begin_nsec: int 

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

276 (inclusive). 

277 """ 

278 

279 end_nsec: int 

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

281 (exclusive). 

282 """ 

283 

284 @cached_property 

285 def value(self) -> Timespan: 

286 """The wrapped value.""" 

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

288 

289 @classmethod 

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

291 """Construct from the wrapped value. 

292 

293 Parameters 

294 ---------- 

295 value : `~lsst.daf.butler.Timespan` 

296 Value to wrap. 

297 

298 Returns 

299 ------- 

300 expression : `TimespanColumnLiteral` 

301 Literal expression object. 

302 """ 

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

304 

305 def __str__(self) -> str: 

306 return str(self.value) 

307 

308 

309@final 

310class RegionColumnLiteral(ColumnLiteralBase): 

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

312 

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

314 first use. 

315 """ 

316 

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

318 

319 encoded: bytes 

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

321 

322 @cached_property 

323 def value(self) -> bytes: 

324 """The wrapped value.""" 

325 return lsst.sphgeom.Region.decode(b64decode(self.encoded)) 

326 

327 @classmethod 

328 def from_value(cls, value: lsst.sphgeom.Region) -> RegionColumnLiteral: 

329 """Construct from the wrapped value. 

330 

331 Parameters 

332 ---------- 

333 value : `~lsst.sphgeom.Region` 

334 Value to wrap. 

335 

336 Returns 

337 ------- 

338 expression : `RegionColumnLiteral` 

339 Literal expression object. 

340 """ 

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

342 

343 def __str__(self) -> str: 

344 return "(region)" 

345 

346 

347ColumnLiteral: TypeAlias = Union[ 

348 IntColumnLiteral, 

349 StringColumnLiteral, 

350 FloatColumnLiteral, 

351 HashColumnLiteral, 

352 UUIDColumnLiteral, 

353 DateTimeColumnLiteral, 

354 TimespanColumnLiteral, 

355 RegionColumnLiteral, 

356] 

357 

358 

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

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

361 

362 Parameters 

363 ---------- 

364 value : `LiteralValue` 

365 Value to wrap. 

366 

367 Returns 

368 ------- 

369 expression : `ColumnLiteral` 

370 Literal expression object. 

371 """ 

372 match value: 

373 case int(): 

374 return IntColumnLiteral.from_value(value) 

375 case str(): 

376 return StringColumnLiteral.from_value(value) 

377 case float(): 

378 return FloatColumnLiteral.from_value(value) 

379 case uuid.UUID(): 

380 return UUIDColumnLiteral.from_value(value) 

381 case bytes(): 

382 return HashColumnLiteral.from_value(value) 

383 case astropy.time.Time(): 

384 return DateTimeColumnLiteral.from_value(value) 

385 case datetime.date(): 

386 return DateTimeColumnLiteral.from_value(astropy.time.Time(value)) 

387 case Timespan(): 

388 return TimespanColumnLiteral.from_value(value) 

389 case lsst.sphgeom.Region(): 

390 return RegionColumnLiteral.from_value(value) 

391 case lsst.sphgeom.LonLat(): 

392 return _make_region_literal_from_lonlat(value) 

393 case astropy.coordinates.SkyCoord(): 

394 icrs = value.transform_to("icrs") 

395 if not icrs.isscalar: 

396 raise ValueError( 

397 "Astropy SkyCoord contained an array of points," 

398 f" but it should be only a single point: {value}" 

399 ) 

400 

401 ra = icrs.ra.degree 

402 dec = icrs.dec.degree 

403 lon_lat = lsst.sphgeom.LonLat.fromDegrees(ra, dec) 

404 return _make_region_literal_from_lonlat(lon_lat) 

405 case numbers.Integral(): 

406 # numpy.int64 and other integer-like values. 

407 return IntColumnLiteral.from_value(int(value)) 

408 

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

410 

411 

412def _make_region_literal_from_lonlat(point: lsst.sphgeom.LonLat) -> RegionColumnLiteral: 

413 vec = lsst.sphgeom.UnitVector3d(point) 

414 # Convert the point to a Region by representing it as a zero-radius 

415 # Circle. 

416 region = lsst.sphgeom.Circle(vec) 

417 return RegionColumnLiteral.from_value(region)