Coverage for python/lsst/daf/butler/queries/visitors.py: 68%
87 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-11 03:16 -0700
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-11 03:16 -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/>.
28from __future__ import annotations
30__all__ = (
31 "ColumnExpressionVisitor",
32 "PredicateVisitor",
33 "SimplePredicateVisitor",
34 "PredicateVisitFlags",
35)
37import enum
38from abc import abstractmethod
39from typing import Generic, TypeVar, final
41from . import tree
43_T = TypeVar("_T")
44_L = TypeVar("_L")
45_A = TypeVar("_A")
46_O = TypeVar("_O")
49class PredicateVisitFlags(enum.Flag):
50 """Flags that provide information about the location of a predicate term
51 in the larger tree.
52 """
54 HAS_AND_SIBLINGS = enum.auto()
55 HAS_OR_SIBLINGS = enum.auto()
56 INVERTED = enum.auto()
59class ColumnExpressionVisitor(Generic[_T]):
60 """A visitor interface for traversing a `ColumnExpression` tree.
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.
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 """
74 @abstractmethod
75 def visit_literal(self, expression: tree.ColumnLiteral) -> _T:
76 """Visit a column expression that wraps a literal value.
78 Parameters
79 ----------
80 expression : `tree.ColumnLiteral`
81 Expression to visit.
83 Returns
84 -------
85 result : `object`
86 Implementation-defined.
87 """
88 raise NotImplementedError()
90 @abstractmethod
91 def visit_dimension_key_reference(self, expression: tree.DimensionKeyReference) -> _T:
92 """Visit a column expression that represents a dimension column.
94 Parameters
95 ----------
96 expression : `tree.DimensionKeyReference`
97 Expression to visit.
99 Returns
100 -------
101 result : `object`
102 Implementation-defined.
103 """
104 raise NotImplementedError()
106 @abstractmethod
107 def visit_dimension_field_reference(self, expression: tree.DimensionFieldReference) -> _T:
108 """Visit a column expression that represents a dimension record field.
110 Parameters
111 ----------
112 expression : `tree.DimensionFieldReference`
113 Expression to visit.
115 Returns
116 -------
117 result : `object`
118 Implementation-defined.
119 """
120 raise NotImplementedError()
122 @abstractmethod
123 def visit_dataset_field_reference(self, expression: tree.DatasetFieldReference) -> _T:
124 """Visit a column expression that represents a dataset field.
126 Parameters
127 ----------
128 expression : `tree.DatasetFieldReference`
129 Expression to visit.
131 Returns
132 -------
133 result : `object`
134 Implementation-defined.
135 """
136 raise NotImplementedError()
138 @abstractmethod
139 def visit_unary_expression(self, expression: tree.UnaryExpression) -> _T:
140 """Visit a column expression that represents a unary operation.
142 Parameters
143 ----------
144 expression : `tree.UnaryExpression`
145 Expression to visit.
147 Returns
148 -------
149 result : `object`
150 Implementation-defined.
151 """
152 raise NotImplementedError()
154 @abstractmethod
155 def visit_binary_expression(self, expression: tree.BinaryExpression) -> _T:
156 """Visit a column expression that wraps a binary operation.
158 Parameters
159 ----------
160 expression : `tree.BinaryExpression`
161 Expression to visit.
163 Returns
164 -------
165 result : `object`
166 Implementation-defined.
167 """
168 raise NotImplementedError()
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.
175 Parameters
176 ----------
177 expression : `tree.Reversed`
178 Expression to visit.
180 Returns
181 -------
182 result : `object`
183 Implementation-defined.
184 """
185 raise NotImplementedError()
188class PredicateVisitor(Generic[_A, _O, _L]):
189 """A visitor interface for traversing a `Predicate`.
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 """
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.
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.
223 Returns
224 -------
225 result : `object`
226 Implementation-defined.
227 """
228 raise NotImplementedError()
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.
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.
243 Returns
244 -------
245 result : `object`
246 Implementation-defined.
247 """
248 raise NotImplementedError()
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.
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.
270 Returns
271 -------
272 result : `object`
273 Implementation-defined.
274 """
275 raise NotImplementedError()
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.
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.
303 Returns
304 -------
305 result : `object`
306 Implementation-defined.
307 """
308 raise NotImplementedError()
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.
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.
333 Returns
334 -------
335 result : `object`
336 Implementation-defined.
337 """
338 raise NotImplementedError()
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.
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.
355 Returns
356 -------
357 result : `object`
358 Implementation-defined.
359 """
360 raise NotImplementedError()
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.
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.
383 Returns
384 -------
385 result : `object`
386 Implementation-defined.
387 """
388 raise NotImplementedError()
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.
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.
403 Returns
404 -------
405 result : `object`
406 Implementation-defined.
407 """
408 raise NotImplementedError()
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 )
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 )
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 )
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`.
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 """
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
461 def visit_is_null(
462 self, operand: tree.ColumnExpression, flags: PredicateVisitFlags
463 ) -> tree.Predicate | None:
464 # Docstring inherited.
465 return None
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
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
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
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
505 return tree.Predicate._from_leaf(original).logical_not()
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
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 )
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
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 )