Coverage for python/lsst/daf/butler/queries/visitors.py: 68%

87 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-05 02:53 -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 "ColumnExpressionVisitor", 

32 "PredicateVisitor", 

33 "SimplePredicateVisitor", 

34 "PredicateVisitFlags", 

35) 

36 

37import enum 

38from abc import abstractmethod 

39from typing import Generic, TypeVar, final 

40 

41from . import tree 

42 

43_T = TypeVar("_T") 

44_L = TypeVar("_L") 

45_A = TypeVar("_A") 

46_O = TypeVar("_O") 

47 

48 

49class PredicateVisitFlags(enum.Flag): 

50 """Flags that provide information about the location of a predicate term 

51 in the larger tree. 

52 """ 

53 

54 HAS_AND_SIBLINGS = enum.auto() 

55 HAS_OR_SIBLINGS = enum.auto() 

56 INVERTED = enum.auto() 

57 

58 

59class ColumnExpressionVisitor(Generic[_T]): 

60 """A visitor interface for traversing a `ColumnExpression` tree. 

61 

62 Notes 

63 ----- 

64 Unlike `Predicate`, the concrete column expression types need to be 

65 public for various reasons, and hence the visitor interface uses them 

66 directly in its arguments. 

67 

68 This interface includes `Reversed` (which is part of the `OrderExpression` 

69 union but not the `ColumnExpression` union) because it is simpler to have 

70 just one visitor interface and disable support for it at runtime as 

71 appropriate. 

72 """ 

73 

74 @abstractmethod 

75 def visit_literal(self, expression: tree.ColumnLiteral) -> _T: 

76 """Visit a column expression that wraps a literal value. 

77 

78 Parameters 

79 ---------- 

80 expression : `tree.ColumnLiteral` 

81 Expression to visit. 

82 

83 Returns 

84 ------- 

85 result : `object` 

86 Implementation-defined. 

87 """ 

88 raise NotImplementedError() 

89 

90 @abstractmethod 

91 def visit_dimension_key_reference(self, expression: tree.DimensionKeyReference) -> _T: 

92 """Visit a column expression that represents a dimension column. 

93 

94 Parameters 

95 ---------- 

96 expression : `tree.DimensionKeyReference` 

97 Expression to visit. 

98 

99 Returns 

100 ------- 

101 result : `object` 

102 Implementation-defined. 

103 """ 

104 raise NotImplementedError() 

105 

106 @abstractmethod 

107 def visit_dimension_field_reference(self, expression: tree.DimensionFieldReference) -> _T: 

108 """Visit a column expression that represents a dimension record field. 

109 

110 Parameters 

111 ---------- 

112 expression : `tree.DimensionFieldReference` 

113 Expression to visit. 

114 

115 Returns 

116 ------- 

117 result : `object` 

118 Implementation-defined. 

119 """ 

120 raise NotImplementedError() 

121 

122 @abstractmethod 

123 def visit_dataset_field_reference(self, expression: tree.DatasetFieldReference) -> _T: 

124 """Visit a column expression that represents a dataset field. 

125 

126 Parameters 

127 ---------- 

128 expression : `tree.DatasetFieldReference` 

129 Expression to visit. 

130 

131 Returns 

132 ------- 

133 result : `object` 

134 Implementation-defined. 

135 """ 

136 raise NotImplementedError() 

137 

138 @abstractmethod 

139 def visit_unary_expression(self, expression: tree.UnaryExpression) -> _T: 

140 """Visit a column expression that represents a unary operation. 

141 

142 Parameters 

143 ---------- 

144 expression : `tree.UnaryExpression` 

145 Expression to visit. 

146 

147 Returns 

148 ------- 

149 result : `object` 

150 Implementation-defined. 

151 """ 

152 raise NotImplementedError() 

153 

154 @abstractmethod 

155 def visit_binary_expression(self, expression: tree.BinaryExpression) -> _T: 

156 """Visit a column expression that wraps a binary operation. 

157 

158 Parameters 

159 ---------- 

160 expression : `tree.BinaryExpression` 

161 Expression to visit. 

162 

163 Returns 

164 ------- 

165 result : `object` 

166 Implementation-defined. 

167 """ 

168 raise NotImplementedError() 

169 

170 @abstractmethod 

171 def visit_reversed(self, expression: tree.Reversed) -> _T: 

172 """Visit a column expression that switches sort order from ascending 

173 to descending. 

174 

175 Parameters 

176 ---------- 

177 expression : `tree.Reversed` 

178 Expression to visit. 

179 

180 Returns 

181 ------- 

182 result : `object` 

183 Implementation-defined. 

184 """ 

185 raise NotImplementedError() 

186 

187 

188class PredicateVisitor(Generic[_A, _O, _L]): 

189 """A visitor interface for traversing a `Predicate`. 

190 

191 Notes 

192 ----- 

193 The concrete `PredicateLeaf` types are only semi-public (they appear in 

194 the serialized form of a `Predicate`, but their types should not generally 

195 be referenced directly outside of the module in which they are defined. 

196 As a result, visiting these objects unpacks their attributes into the 

197 visit method arguments. 

198 """ 

199 

200 @abstractmethod 

201 def visit_comparison( 

202 self, 

203 a: tree.ColumnExpression, 

204 operator: tree.ComparisonOperator, 

205 b: tree.ColumnExpression, 

206 flags: PredicateVisitFlags, 

207 ) -> _L: 

208 """Visit a binary comparison between column expressions. 

209 

210 Parameters 

211 ---------- 

212 a : `tree.ColumnExpression` 

213 First column expression in the comparison. 

214 operator : `str` 

215 Enumerated string representing the comparison operator to apply. 

216 May be and of "==", "!=", "<", ">", "<=", ">=", or "overlaps". 

217 b : `tree.ColumnExpression` 

218 Second column expression in the comparison. 

219 flags : `PredicateVisitFlags` 

220 Information about where this leaf appears in the larger predicate 

221 tree. 

222 

223 Returns 

224 ------- 

225 result : `object` 

226 Implementation-defined. 

227 """ 

228 raise NotImplementedError() 

229 

230 @abstractmethod 

231 def visit_is_null(self, operand: tree.ColumnExpression, flags: PredicateVisitFlags) -> _L: 

232 """Visit a predicate leaf that tests whether a column expression is 

233 NULL. 

234 

235 Parameters 

236 ---------- 

237 operand : `tree.ColumnExpression` 

238 Column expression to test. 

239 flags : `PredicateVisitFlags` 

240 Information about where this leaf appears in the larger predicate 

241 tree. 

242 

243 Returns 

244 ------- 

245 result : `object` 

246 Implementation-defined. 

247 """ 

248 raise NotImplementedError() 

249 

250 @abstractmethod 

251 def visit_in_container( 

252 self, 

253 member: tree.ColumnExpression, 

254 container: tuple[tree.ColumnExpression, ...], 

255 flags: PredicateVisitFlags, 

256 ) -> _L: 

257 """Visit a predicate leaf that tests whether a column expression is 

258 a member of a container. 

259 

260 Parameters 

261 ---------- 

262 member : `tree.ColumnExpression` 

263 Column expression that may be a member of the container. 

264 container : `~collections.abc.Iterable` [ `tree.ColumnExpression` ] 

265 Container of column expressions to test for membership in. 

266 flags : `PredicateVisitFlags` 

267 Information about where this leaf appears in the larger predicate 

268 tree. 

269 

270 Returns 

271 ------- 

272 result : `object` 

273 Implementation-defined. 

274 """ 

275 raise NotImplementedError() 

276 

277 @abstractmethod 

278 def visit_in_range( 

279 self, 

280 member: tree.ColumnExpression, 

281 start: int, 

282 stop: int | None, 

283 step: int, 

284 flags: PredicateVisitFlags, 

285 ) -> _L: 

286 """Visit a predicate leaf that tests whether a column expression is 

287 a member of an integer range. 

288 

289 Parameters 

290 ---------- 

291 member : `tree.ColumnExpression` 

292 Column expression that may be a member of the range. 

293 start : `int`, optional 

294 Beginning of the range, inclusive. 

295 stop : `int` or `None`, optional 

296 End of the range, exclusive. 

297 step : `int`, optional 

298 Offset between values in the range. 

299 flags : `PredicateVisitFlags` 

300 Information about where this leaf appears in the larger predicate 

301 tree. 

302 

303 Returns 

304 ------- 

305 result : `object` 

306 Implementation-defined. 

307 """ 

308 raise NotImplementedError() 

309 

310 @abstractmethod 

311 def visit_in_query_tree( 

312 self, 

313 member: tree.ColumnExpression, 

314 column: tree.ColumnExpression, 

315 query_tree: tree.QueryTree, 

316 flags: PredicateVisitFlags, 

317 ) -> _L: 

318 """Visit a predicate leaf that tests whether a column expression is 

319 a member of a container. 

320 

321 Parameters 

322 ---------- 

323 member : `tree.ColumnExpression` 

324 Column expression that may be present in the query. 

325 column : `tree.ColumnExpression` 

326 Column to project from the query. 

327 query_tree : `QueryTree` 

328 Query tree to select from. 

329 flags : `PredicateVisitFlags` 

330 Information about where this leaf appears in the larger predicate 

331 tree. 

332 

333 Returns 

334 ------- 

335 result : `object` 

336 Implementation-defined. 

337 """ 

338 raise NotImplementedError() 

339 

340 @abstractmethod 

341 def apply_logical_not(self, original: tree.PredicateLeaf, result: _L, flags: PredicateVisitFlags) -> _L: 

342 """Apply a logical NOT to the result of visiting an inverted predicate 

343 leaf. 

344 

345 Parameters 

346 ---------- 

347 original : `PredicateLeaf` 

348 The original operand of the logical NOT operation. 

349 result : `object` 

350 Implementation-defined result of visiting the operand. 

351 flags : `PredicateVisitFlags` 

352 Information about where this leaf appears in the larger predicate 

353 tree. Never has `PredicateVisitFlags.INVERTED` set. 

354 

355 Returns 

356 ------- 

357 result : `object` 

358 Implementation-defined. 

359 """ 

360 raise NotImplementedError() 

361 

362 @abstractmethod 

363 def apply_logical_or( 

364 self, 

365 originals: tuple[tree.PredicateLeaf, ...], 

366 results: tuple[_L, ...], 

367 flags: PredicateVisitFlags, 

368 ) -> _O: 

369 """Apply a logical OR operation to the result of visiting a `tuple` of 

370 predicate leaf objects. 

371 

372 Parameters 

373 ---------- 

374 originals : `tuple` [ `PredicateLeaf`, ... ] 

375 Original leaf objects in the logical OR. 

376 results : `tuple` [ `object`, ... ] 

377 Result of visiting the leaf objects. 

378 flags : `PredicateVisitFlags` 

379 Information about where this leaf appears in the larger predicate 

380 tree. Never has `PredicateVisitFlags.INVERTED` or 

381 `PredicateVisitFlags.HAS_OR_SIBLINGS` set. 

382 

383 Returns 

384 ------- 

385 result : `object` 

386 Implementation-defined. 

387 """ 

388 raise NotImplementedError() 

389 

390 @abstractmethod 

391 def apply_logical_and(self, originals: tree.PredicateOperands, results: tuple[_O, ...]) -> _A: 

392 """Apply a logical AND operation to the result of visiting a nested 

393 `tuple` of predicate leaf objects. 

394 

395 Parameters 

396 ---------- 

397 originals : `tuple` [ `tuple` [ `PredicateLeaf`, ... ], ... ] 

398 Nested tuple of predicate leaf objects, with inner tuples 

399 corresponding to groups that should be combined with logical OR. 

400 results : `tuple` [ `object`, ... ] 

401 Result of visiting the leaf objects. 

402 

403 Returns 

404 ------- 

405 result : `object` 

406 Implementation-defined. 

407 """ 

408 raise NotImplementedError() 

409 

410 @final 

411 def _visit_logical_not(self, operand: tree.LogicalNotOperand, flags: PredicateVisitFlags) -> _L: 

412 return self.apply_logical_not( 

413 operand, operand.visit(self, flags | PredicateVisitFlags.INVERTED), flags 

414 ) 

415 

416 @final 

417 def _visit_logical_or(self, operands: tuple[tree.PredicateLeaf, ...], flags: PredicateVisitFlags) -> _O: 

418 nested_flags = flags 

419 if len(operands) > 1: 

420 nested_flags |= PredicateVisitFlags.HAS_OR_SIBLINGS 

421 return self.apply_logical_or( 

422 operands, tuple([operand.visit(self, nested_flags) for operand in operands]), flags 

423 ) 

424 

425 @final 

426 def _visit_logical_and(self, operands: tree.PredicateOperands) -> _A: 

427 if len(operands) > 1: 

428 nested_flags = PredicateVisitFlags.HAS_AND_SIBLINGS 

429 else: 

430 nested_flags = PredicateVisitFlags(0) 

431 return self.apply_logical_and( 

432 operands, tuple([self._visit_logical_or(or_group, nested_flags) for or_group in operands]) 

433 ) 

434 

435 

436class SimplePredicateVisitor( 

437 PredicateVisitor[tree.Predicate | None, tree.Predicate | None, tree.Predicate | None] 

438): 

439 """An intermediate base class for predicate visitor implementations that 

440 either return `None` or a new `Predicate`. 

441 

442 Notes 

443 ----- 

444 This class implements all leaf-node visitation methods to return `None`, 

445 which is interpreted by the ``apply*`` method implementations as indicating 

446 that the leaf is unmodified. Subclasses can thus override only certain 

447 visitation methods and either return `None` if there is no result, or 

448 return a replacement `Predicate` to construct a new tree. 

449 """ 

450 

451 def visit_comparison( 

452 self, 

453 a: tree.ColumnExpression, 

454 operator: tree.ComparisonOperator, 

455 b: tree.ColumnExpression, 

456 flags: PredicateVisitFlags, 

457 ) -> tree.Predicate | None: 

458 # Docstring inherited. 

459 return None 

460 

461 def visit_is_null( 

462 self, operand: tree.ColumnExpression, flags: PredicateVisitFlags 

463 ) -> tree.Predicate | None: 

464 # Docstring inherited. 

465 return None 

466 

467 def visit_in_container( 

468 self, 

469 member: tree.ColumnExpression, 

470 container: tuple[tree.ColumnExpression, ...], 

471 flags: PredicateVisitFlags, 

472 ) -> tree.Predicate | None: 

473 # Docstring inherited. 

474 return None 

475 

476 def visit_in_range( 

477 self, 

478 member: tree.ColumnExpression, 

479 start: int, 

480 stop: int | None, 

481 step: int, 

482 flags: PredicateVisitFlags, 

483 ) -> tree.Predicate | None: 

484 # Docstring inherited. 

485 return None 

486 

487 def visit_in_query_tree( 

488 self, 

489 member: tree.ColumnExpression, 

490 column: tree.ColumnExpression, 

491 query_tree: tree.QueryTree, 

492 flags: PredicateVisitFlags, 

493 ) -> tree.Predicate | None: 

494 # Docstring inherited. 

495 return None 

496 

497 def apply_logical_not( 

498 self, original: tree.PredicateLeaf, result: tree.Predicate | None, flags: PredicateVisitFlags 

499 ) -> tree.Predicate | None: 

500 # Docstring inherited. 

501 if result is None: 

502 return None 

503 from . import tree 

504 

505 return tree.Predicate._from_leaf(original).logical_not() 

506 

507 def apply_logical_or( 

508 self, 

509 originals: tuple[tree.PredicateLeaf, ...], 

510 results: tuple[tree.Predicate | None, ...], 

511 flags: PredicateVisitFlags, 

512 ) -> tree.Predicate | None: 

513 # Docstring inherited. 

514 if all(result is None for result in results): 

515 return None 

516 from . import tree 

517 

518 return tree.Predicate.from_bool(False).logical_or( 

519 *[ 

520 tree.Predicate._from_leaf(original) if result is None else result 

521 for original, result in zip(originals, results) 

522 ] 

523 ) 

524 

525 def apply_logical_and( 

526 self, 

527 originals: tree.PredicateOperands, 

528 results: tuple[tree.Predicate | None, ...], 

529 ) -> tree.Predicate | None: 

530 # Docstring inherited. 

531 if all(result is None for result in results): 

532 return None 

533 from . import tree 

534 

535 return tree.Predicate.from_bool(True).logical_and( 

536 *[ 

537 tree.Predicate._from_or_group(original) if result is None else result 

538 for original, result in zip(originals, results) 

539 ] 

540 )