Coverage for python/lsst/daf/butler/queries/tree/_column_expression.py: 48%
110 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-07 02:46 -0700
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-07 02:46 -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 "ColumnExpression",
32 "OrderExpression",
33 "UnaryExpression",
34 "BinaryExpression",
35 "Reversed",
36 "UnaryOperator",
37 "BinaryOperator",
38 "validate_order_expression",
39)
41from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias, TypeVar, final
43import pydantic
45from ..._exceptions import InvalidQueryError
46from ...column_spec import ColumnType
47from ._base import ColumnExpressionBase
48from ._column_literal import ColumnLiteral
49from ._column_reference import ColumnReference
50from ._column_set import ColumnSet
52if TYPE_CHECKING:
53 from ..visitors import ColumnExpressionVisitor
56_T = TypeVar("_T")
59UnaryOperator: TypeAlias = Literal["-", "begin_of", "end_of"]
60BinaryOperator: TypeAlias = Literal["+", "-", "*", "/", "%"]
63@final
64class UnaryExpression(ColumnExpressionBase):
65 """A unary operation on a column expression that returns a non-boolean
66 value.
67 """
69 expression_type: Literal["unary"] = "unary"
71 operand: ColumnExpression
72 """Expression this one operates on."""
74 operator: UnaryOperator
75 """Operator this expression applies."""
77 def gather_required_columns(self, columns: ColumnSet) -> None:
78 # Docstring inherited.
79 self.operand.gather_required_columns(columns)
81 @property
82 def column_type(self) -> ColumnType:
83 # Docstring inherited.
84 match self.operator:
85 case "-":
86 return self.operand.column_type
87 case "begin_of" | "end_of":
88 return "datetime"
89 raise AssertionError(f"Invalid unary expression operator {self.operator}.")
91 def __str__(self) -> str:
92 s = str(self.operand)
93 if not (self.operand.is_literal or self.operand.is_column_reference):
94 s = f"({s})"
95 match self.operator:
96 case "-":
97 return f"-{s}"
98 case "begin_of":
99 return f"{s}.begin"
100 case "end_of":
101 return f"{s}.end"
103 @pydantic.model_validator(mode="after")
104 def _validate_types(self) -> UnaryExpression:
105 match (self.operator, self.operand.column_type):
106 case ("-", "int" | "float"):
107 pass
108 case ("begin_of" | "end_of", "timespan"):
109 pass
110 case _:
111 raise InvalidQueryError(
112 f"Invalid column type {self.operand.column_type} for operator {self.operator!r}."
113 )
114 return self
116 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
117 # Docstring inherited.
118 return visitor.visit_unary_expression(self)
121@final
122class BinaryExpression(ColumnExpressionBase):
123 """A binary operation on column expressions that returns a non-boolean
124 value.
125 """
127 expression_type: Literal["binary"] = "binary"
129 a: ColumnExpression
130 """Left-hand side expression this one operates on."""
132 b: ColumnExpression
133 """Right-hand side expression this one operates on."""
135 operator: BinaryOperator
136 """Operator this expression applies.
138 Integer '/' and '%' are defined as in SQL, not Python (though the
139 definitions are the same for positive arguments).
140 """
142 def gather_required_columns(self, columns: ColumnSet) -> None:
143 # Docstring inherited.
144 self.a.gather_required_columns(columns)
145 self.b.gather_required_columns(columns)
147 @property
148 def column_type(self) -> ColumnType:
149 # Docstring inherited.
150 return self.a.column_type
152 def __str__(self) -> str:
153 a = str(self.a)
154 b = str(self.b)
155 if not (self.a.is_literal or self.a.is_column_reference):
156 a = f"({a})"
157 if not (self.b.is_literal or self.b.is_column_reference):
158 b = f"({b})"
159 return f"{a} {self.operator} {b}"
161 @pydantic.model_validator(mode="after")
162 def _validate_types(self) -> BinaryExpression:
163 if self.a.column_type != self.b.column_type:
164 raise InvalidQueryError(
165 f"Column types for operator {self.operator} do not agree "
166 f"({self.a.column_type}, {self.b.column_type})."
167 )
168 match (self.operator, self.a.column_type):
169 case ("+" | "-" | "*" | "/", "int" | "float"):
170 pass
171 case ("%", "int"):
172 pass
173 case _:
174 raise InvalidQueryError(
175 f"Invalid column type {self.a.column_type} for operator {self.operator!r}."
176 )
177 return self
179 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
180 # Docstring inherited.
181 return visitor.visit_binary_expression(self)
184# Union without Pydantic annotation for the discriminator, for use in nesting
185# in other unions that will add that annotation. It's not clear whether it
186# would work to just nest the annotated ones, but it seems safest not to rely
187# on undocumented behavior.
188_ColumnExpression: TypeAlias = ColumnLiteral | ColumnReference | UnaryExpression | BinaryExpression
191ColumnExpression: TypeAlias = Annotated[_ColumnExpression, pydantic.Field(discriminator="expression_type")]
194@final
195class Reversed(ColumnExpressionBase):
196 """A tag wrapper for `ColumnExpression` that indicates sorting in
197 reverse order.
198 """
200 expression_type: Literal["reversed"] = "reversed"
202 operand: ColumnExpression
203 """Expression to sort on in reverse."""
205 def gather_required_columns(self, columns: ColumnSet) -> None:
206 # Docstring inherited.
207 self.operand.gather_required_columns(columns)
209 @property
210 def column_type(self) -> ColumnType:
211 # Docstring inherited.
212 return self.operand.column_type
214 def __str__(self) -> str:
215 return f"{self.operand} DESC"
217 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
218 # Docstring inherited.
219 return visitor.visit_reversed(self)
222def validate_order_expression(expression: _ColumnExpression | Reversed) -> _ColumnExpression | Reversed:
223 """Check that a column expression can be used for sorting.
225 Parameters
226 ----------
227 expression : `OrderExpression`
228 Expression to check.
230 Returns
231 -------
232 expression : `OrderExpression`
233 The checked expression; returned to make this usable as a Pydantic
234 validator.
236 Raises
237 ------
238 InvalidQueryError
239 Raised if this expression is not one that can be used for sorting.
240 """
241 if expression.column_type not in ("int", "string", "float", "datetime"):
242 raise InvalidQueryError(f"Column type {expression.column_type} of {expression} is not ordered.")
243 return expression
246OrderExpression: TypeAlias = Annotated[
247 _ColumnExpression | Reversed,
248 pydantic.Field(discriminator="expression_type"),
249 pydantic.AfterValidator(validate_order_expression),
250]