Coverage for python/lsst/daf/butler/queries/tree/_column_expression.py: 46%
126 statements
« prev ^ index » next coverage.py v7.4.3, created at 2024-03-05 11:36 +0000
« prev ^ index » next coverage.py v7.4.3, created at 2024-03-05 11:36 +0000
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, Union, final
43import pydantic
45from ...column_spec import ColumnType
46from ._base import ColumnExpressionBase, InvalidQueryError
47from ._column_literal import ColumnLiteral
48from ._column_reference import ColumnReference
49from ._column_set import ColumnSet
51if TYPE_CHECKING:
52 from ..visitors import ColumnExpressionVisitor
55_T = TypeVar("_T")
58UnaryOperator: TypeAlias = Literal["-", "begin_of", "end_of"]
59BinaryOperator: TypeAlias = Literal["+", "-", "*", "/", "%"]
62@final
63class UnaryExpression(ColumnExpressionBase):
64 """A unary operation on a column expression that returns a non-boolean
65 value.
66 """
68 expression_type: Literal["unary"] = "unary"
70 operand: ColumnExpression
71 """Expression this one operates on."""
73 operator: UnaryOperator
74 """Operator this expression applies."""
76 def gather_required_columns(self, columns: ColumnSet) -> None:
77 # Docstring inherited.
78 self.operand.gather_required_columns(columns)
80 @property
81 def precedence(self) -> int:
82 # Docstring inherited.
83 return 1
85 @property
86 def column_type(self) -> ColumnType:
87 # Docstring inherited.
88 match self.operator:
89 case "-":
90 return self.operand.column_type
91 case "begin_of" | "end_of":
92 return "datetime"
93 raise AssertionError(f"Invalid unary expression operator {self.operator}.")
95 def __str__(self) -> str:
96 s = str(self.operand)
97 match self.operator:
98 case "-":
99 return f"-{s}"
100 case "begin_of":
101 return f"{s}.begin"
102 case "end_of":
103 return f"{s}.end"
105 @pydantic.model_validator(mode="after")
106 def _validate_types(self) -> UnaryExpression:
107 match (self.operator, self.operand.column_type):
108 case ("-", "int" | "float"):
109 pass
110 case ("begin_of" | "end_of", "timespan"):
111 pass
112 case _:
113 raise InvalidQueryError(
114 f"Invalid column type {self.operand.column_type} for operator {self.operator!r}."
115 )
116 return self
118 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
119 # Docstring inherited.
120 return visitor.visit_unary_expression(self)
123@final
124class BinaryExpression(ColumnExpressionBase):
125 """A binary operation on column expressions that returns a non-boolean
126 value.
127 """
129 expression_type: Literal["binary"] = "binary"
131 a: ColumnExpression
132 """Left-hand side expression this one operates on."""
134 b: ColumnExpression
135 """Right-hand side expression this one operates on."""
137 operator: BinaryOperator
138 """Operator this expression applies.
140 Integer '/' and '%' are defined as in SQL, not Python (though the
141 definitions are the same for positive arguments).
142 """
144 def gather_required_columns(self, columns: ColumnSet) -> None:
145 # Docstring inherited.
146 self.a.gather_required_columns(columns)
147 self.b.gather_required_columns(columns)
149 @property
150 def precedence(self) -> int:
151 # Docstring inherited.
152 match self.operator:
153 case "*" | "/" | "%":
154 return 2
155 case "+" | "-":
156 return 3
158 @property
159 def column_type(self) -> ColumnType:
160 # Docstring inherited.
161 return self.a.column_type
163 def __str__(self) -> str:
164 a = str(self.a)
165 b = str(self.b)
166 match self.operator:
167 case "*" | "+":
168 if self.a.precedence > self.precedence:
169 a = f"({a})"
170 if self.b.precedence > self.precedence:
171 b = f"({b})"
172 case _:
173 if self.a.precedence >= self.precedence:
174 a = f"({a})"
175 if self.b.precedence >= self.precedence:
176 b = f"({b})"
177 return f"{a} {self.operator} {b}"
179 @pydantic.model_validator(mode="after")
180 def _validate_types(self) -> BinaryExpression:
181 if self.a.column_type != self.b.column_type:
182 raise InvalidQueryError(
183 f"Column types for operator {self.operator} do not agree "
184 f"({self.a.column_type}, {self.b.column_type})."
185 )
186 match (self.operator, self.a.column_type):
187 case ("+" | "-" | "*" | "/", "int" | "float"):
188 pass
189 case ("%", "int"):
190 pass
191 case _:
192 raise InvalidQueryError(
193 f"Invalid column type {self.a.column_type} for operator {self.operator!r}."
194 )
195 return self
197 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
198 # Docstring inherited.
199 return visitor.visit_binary_expression(self)
202# Union without Pydantic annotation for the discriminator, for use in nesting
203# in other unions that will add that annotation. It's not clear whether it
204# would work to just nest the annotated ones, but it seems safest not to rely
205# on undocumented behavior.
206_ColumnExpression: TypeAlias = Union[
207 ColumnLiteral,
208 ColumnReference,
209 UnaryExpression,
210 BinaryExpression,
211]
214ColumnExpression: TypeAlias = Annotated[_ColumnExpression, pydantic.Field(discriminator="expression_type")]
217@final
218class Reversed(ColumnExpressionBase):
219 """A tag wrapper for `ColumnExpression` that indicates sorting in
220 reverse order.
221 """
223 expression_type: Literal["reversed"] = "reversed"
225 operand: ColumnExpression
226 """Expression to sort on in reverse."""
228 def gather_required_columns(self, columns: ColumnSet) -> None:
229 # Docstring inherited.
230 self.operand.gather_required_columns(columns)
232 @property
233 def precedence(self) -> int:
234 # Docstring inherited.
235 raise AssertionError("Order-reversed expressions can never be nested in other column expressions.")
237 @property
238 def column_type(self) -> ColumnType:
239 # Docstring inherited.
240 return self.operand.column_type
242 def __str__(self) -> str:
243 return f"{self.operand} DESC"
245 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
246 # Docstring inherited.
247 return visitor.visit_reversed(self)
250def validate_order_expression(expression: _ColumnExpression | Reversed) -> _ColumnExpression | Reversed:
251 """Check that a column expression can be used for sorting.
253 Parameters
254 ----------
255 expression : `OrderExpression`
256 Expression to check.
258 Returns
259 -------
260 expression : `OrderExpression`
261 The checked expression; returned to make this usable as a Pydantic
262 validator.
264 Raises
265 ------
266 InvalidQueryError
267 Raised if this expression is not one that can be used for sorting.
268 """
269 if expression.column_type not in ("int", "string", "float", "datetime"):
270 raise InvalidQueryError(f"Column type {expression.column_type} of {expression} is not ordered.")
271 return expression
274OrderExpression: TypeAlias = Annotated[
275 Union[_ColumnExpression, Reversed],
276 pydantic.Field(discriminator="expression_type"),
277 pydantic.AfterValidator(validate_order_expression),
278]