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

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 "PredicateVisitFlags", 

33 "PredicateVisitor", 

34 "SimplePredicateVisitor", 

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_boolean_wrapper(self, value: tree.ColumnExpression, flags: PredicateVisitFlags) -> _L: 

202 """Visit a boolean-valued column expression. 

203 

204 Parameters 

205 ---------- 

206 value : `tree.ColumnExpression` 

207 Column expression, guaranteed to have `column_type == "bool"`. 

208 flags : `PredicateVisitFlags` 

209 Information about where this leaf appears in the larger predicate 

210 tree. 

211 

212 Returns 

213 ------- 

214 result : `object` 

215 Implementation-defined. 

216 """ 

217 raise NotImplementedError() 

218 

219 @abstractmethod 

220 def visit_comparison( 

221 self, 

222 a: tree.ColumnExpression, 

223 operator: tree.ComparisonOperator, 

224 b: tree.ColumnExpression, 

225 flags: PredicateVisitFlags, 

226 ) -> _L: 

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

228 

229 Parameters 

230 ---------- 

231 a : `tree.ColumnExpression` 

232 First column expression in the comparison. 

233 operator : `str` 

234 Enumerated string representing the comparison operator to apply. 

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

236 b : `tree.ColumnExpression` 

237 Second column expression in the comparison. 

238 flags : `PredicateVisitFlags` 

239 Information about where this leaf appears in the larger predicate 

240 tree. 

241 

242 Returns 

243 ------- 

244 result : `object` 

245 Implementation-defined. 

246 """ 

247 raise NotImplementedError() 

248 

249 @abstractmethod 

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

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

252 NULL. 

253 

254 Parameters 

255 ---------- 

256 operand : `tree.ColumnExpression` 

257 Column expression to test. 

258 flags : `PredicateVisitFlags` 

259 Information about where this leaf appears in the larger predicate 

260 tree. 

261 

262 Returns 

263 ------- 

264 result : `object` 

265 Implementation-defined. 

266 """ 

267 raise NotImplementedError() 

268 

269 @abstractmethod 

270 def visit_in_container( 

271 self, 

272 member: tree.ColumnExpression, 

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

274 flags: PredicateVisitFlags, 

275 ) -> _L: 

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

277 a member of a container. 

278 

279 Parameters 

280 ---------- 

281 member : `tree.ColumnExpression` 

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

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

284 Container of column expressions to test for membership in. 

285 flags : `PredicateVisitFlags` 

286 Information about where this leaf appears in the larger predicate 

287 tree. 

288 

289 Returns 

290 ------- 

291 result : `object` 

292 Implementation-defined. 

293 """ 

294 raise NotImplementedError() 

295 

296 @abstractmethod 

297 def visit_in_range( 

298 self, 

299 member: tree.ColumnExpression, 

300 start: int, 

301 stop: int | None, 

302 step: int, 

303 flags: PredicateVisitFlags, 

304 ) -> _L: 

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

306 a member of an integer range. 

307 

308 Parameters 

309 ---------- 

310 member : `tree.ColumnExpression` 

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

312 start : `int`, optional 

313 Beginning of the range, inclusive. 

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

315 End of the range, exclusive. 

316 step : `int`, optional 

317 Offset between values in the range. 

318 flags : `PredicateVisitFlags` 

319 Information about where this leaf appears in the larger predicate 

320 tree. 

321 

322 Returns 

323 ------- 

324 result : `object` 

325 Implementation-defined. 

326 """ 

327 raise NotImplementedError() 

328 

329 @abstractmethod 

330 def visit_in_query_tree( 

331 self, 

332 member: tree.ColumnExpression, 

333 column: tree.ColumnExpression, 

334 query_tree: tree.QueryTree, 

335 flags: PredicateVisitFlags, 

336 ) -> _L: 

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

338 a member of a container. 

339 

340 Parameters 

341 ---------- 

342 member : `tree.ColumnExpression` 

343 Column expression that may be present in the query. 

344 column : `tree.ColumnExpression` 

345 Column to project from the query. 

346 query_tree : `QueryTree` 

347 Query tree to select from. 

348 flags : `PredicateVisitFlags` 

349 Information about where this leaf appears in the larger predicate 

350 tree. 

351 

352 Returns 

353 ------- 

354 result : `object` 

355 Implementation-defined. 

356 """ 

357 raise NotImplementedError() 

358 

359 @abstractmethod 

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

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

362 leaf. 

363 

364 Parameters 

365 ---------- 

366 original : `PredicateLeaf` 

367 The original operand of the logical NOT operation. 

368 result : `object` 

369 Implementation-defined result of visiting the operand. 

370 flags : `PredicateVisitFlags` 

371 Information about where this leaf appears in the larger predicate 

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

373 

374 Returns 

375 ------- 

376 result : `object` 

377 Implementation-defined. 

378 """ 

379 raise NotImplementedError() 

380 

381 @abstractmethod 

382 def apply_logical_or( 

383 self, 

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

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

386 flags: PredicateVisitFlags, 

387 ) -> _O: 

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

389 predicate leaf objects. 

390 

391 Parameters 

392 ---------- 

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

394 Original leaf objects in the logical OR. 

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

396 Result of visiting the leaf objects. 

397 flags : `PredicateVisitFlags` 

398 Information about where this leaf appears in the larger predicate 

399 tree. Never has `PredicateVisitFlags.INVERTED` or 

400 `PredicateVisitFlags.HAS_OR_SIBLINGS` set. 

401 

402 Returns 

403 ------- 

404 result : `object` 

405 Implementation-defined. 

406 """ 

407 raise NotImplementedError() 

408 

409 @abstractmethod 

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

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

412 `tuple` of predicate leaf objects. 

413 

414 Parameters 

415 ---------- 

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

417 Nested tuple of predicate leaf objects, with inner tuples 

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

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

420 Result of visiting the leaf objects. 

421 

422 Returns 

423 ------- 

424 result : `object` 

425 Implementation-defined. 

426 """ 

427 raise NotImplementedError() 

428 

429 @final 

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

431 return self.apply_logical_not( 

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

433 ) 

434 

435 @final 

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

437 nested_flags = flags 

438 if len(operands) > 1: 

439 nested_flags |= PredicateVisitFlags.HAS_OR_SIBLINGS 

440 return self.apply_logical_or( 

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

442 ) 

443 

444 @final 

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

446 if len(operands) > 1: 

447 nested_flags = PredicateVisitFlags.HAS_AND_SIBLINGS 

448 else: 

449 nested_flags = PredicateVisitFlags(0) 

450 return self.apply_logical_and( 

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

452 ) 

453 

454 

455class SimplePredicateVisitor( 

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

457): 

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

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

460 

461 Notes 

462 ----- 

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

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

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

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

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

468 """ 

469 

470 def visit_boolean_wrapper( 

471 self, value: tree.ColumnExpression, flags: PredicateVisitFlags 

472 ) -> tree.Predicate | None: 

473 return None 

474 

475 def visit_comparison( 

476 self, 

477 a: tree.ColumnExpression, 

478 operator: tree.ComparisonOperator, 

479 b: tree.ColumnExpression, 

480 flags: PredicateVisitFlags, 

481 ) -> tree.Predicate | None: 

482 # Docstring inherited. 

483 return None 

484 

485 def visit_is_null( 

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

487 ) -> tree.Predicate | None: 

488 # Docstring inherited. 

489 return None 

490 

491 def visit_in_container( 

492 self, 

493 member: tree.ColumnExpression, 

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

495 flags: PredicateVisitFlags, 

496 ) -> tree.Predicate | None: 

497 # Docstring inherited. 

498 return None 

499 

500 def visit_in_range( 

501 self, 

502 member: tree.ColumnExpression, 

503 start: int, 

504 stop: int | None, 

505 step: int, 

506 flags: PredicateVisitFlags, 

507 ) -> tree.Predicate | None: 

508 # Docstring inherited. 

509 return None 

510 

511 def visit_in_query_tree( 

512 self, 

513 member: tree.ColumnExpression, 

514 column: tree.ColumnExpression, 

515 query_tree: tree.QueryTree, 

516 flags: PredicateVisitFlags, 

517 ) -> tree.Predicate | None: 

518 # Docstring inherited. 

519 return None 

520 

521 def apply_logical_not( 

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

523 ) -> tree.Predicate | None: 

524 # Docstring inherited. 

525 if result is None: 

526 return None 

527 from . import tree 

528 

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

530 

531 def apply_logical_or( 

532 self, 

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

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

535 flags: PredicateVisitFlags, 

536 ) -> tree.Predicate | None: 

537 # Docstring inherited. 

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

539 return None 

540 from . import tree 

541 

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

543 *[ 

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

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

546 ] 

547 ) 

548 

549 def apply_logical_and( 

550 self, 

551 originals: tree.PredicateOperands, 

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

553 ) -> tree.Predicate | None: 

554 # Docstring inherited. 

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

556 return None 

557 from . import tree 

558 

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

560 *[ 

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

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

563 ] 

564 )