Coverage for python/lsst/daf/butler/queries/tree/_column_expression.py: 48%
109 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 02:53 -0700
« 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/>.
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 ...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 column_type(self) -> ColumnType:
82 # Docstring inherited.
83 match self.operator:
84 case "-":
85 return self.operand.column_type
86 case "begin_of" | "end_of":
87 return "datetime"
88 raise AssertionError(f"Invalid unary expression operator {self.operator}.")
90 def __str__(self) -> str:
91 s = str(self.operand)
92 if not (self.operand.is_literal or self.operand.is_column_reference):
93 s = f"({s})"
94 match self.operator:
95 case "-":
96 return f"-{s}"
97 case "begin_of":
98 return f"{s}.begin"
99 case "end_of":
100 return f"{s}.end"
102 @pydantic.model_validator(mode="after")
103 def _validate_types(self) -> UnaryExpression:
104 match (self.operator, self.operand.column_type):
105 case ("-", "int" | "float"):
106 pass
107 case ("begin_of" | "end_of", "timespan"):
108 pass
109 case _:
110 raise InvalidQueryError(
111 f"Invalid column type {self.operand.column_type} for operator {self.operator!r}."
112 )
113 return self
115 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
116 # Docstring inherited.
117 return visitor.visit_unary_expression(self)
120@final
121class BinaryExpression(ColumnExpressionBase):
122 """A binary operation on column expressions that returns a non-boolean
123 value.
124 """
126 expression_type: Literal["binary"] = "binary"
128 a: ColumnExpression
129 """Left-hand side expression this one operates on."""
131 b: ColumnExpression
132 """Right-hand side expression this one operates on."""
134 operator: BinaryOperator
135 """Operator this expression applies.
137 Integer '/' and '%' are defined as in SQL, not Python (though the
138 definitions are the same for positive arguments).
139 """
141 def gather_required_columns(self, columns: ColumnSet) -> None:
142 # Docstring inherited.
143 self.a.gather_required_columns(columns)
144 self.b.gather_required_columns(columns)
146 @property
147 def column_type(self) -> ColumnType:
148 # Docstring inherited.
149 return self.a.column_type
151 def __str__(self) -> str:
152 a = str(self.a)
153 b = str(self.b)
154 if not (self.a.is_literal or self.a.is_column_reference):
155 a = f"({a})"
156 if not (self.b.is_literal or self.b.is_column_reference):
157 b = f"({b})"
158 return f"{a} {self.operator} {b}"
160 @pydantic.model_validator(mode="after")
161 def _validate_types(self) -> BinaryExpression:
162 if self.a.column_type != self.b.column_type:
163 raise InvalidQueryError(
164 f"Column types for operator {self.operator} do not agree "
165 f"({self.a.column_type}, {self.b.column_type})."
166 )
167 match (self.operator, self.a.column_type):
168 case ("+" | "-" | "*" | "/", "int" | "float"):
169 pass
170 case ("%", "int"):
171 pass
172 case _:
173 raise InvalidQueryError(
174 f"Invalid column type {self.a.column_type} for operator {self.operator!r}."
175 )
176 return self
178 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
179 # Docstring inherited.
180 return visitor.visit_binary_expression(self)
183# Union without Pydantic annotation for the discriminator, for use in nesting
184# in other unions that will add that annotation. It's not clear whether it
185# would work to just nest the annotated ones, but it seems safest not to rely
186# on undocumented behavior.
187_ColumnExpression: TypeAlias = ColumnLiteral | ColumnReference | UnaryExpression | BinaryExpression
190ColumnExpression: TypeAlias = Annotated[_ColumnExpression, pydantic.Field(discriminator="expression_type")]
193@final
194class Reversed(ColumnExpressionBase):
195 """A tag wrapper for `ColumnExpression` that indicates sorting in
196 reverse order.
197 """
199 expression_type: Literal["reversed"] = "reversed"
201 operand: ColumnExpression
202 """Expression to sort on in reverse."""
204 def gather_required_columns(self, columns: ColumnSet) -> None:
205 # Docstring inherited.
206 self.operand.gather_required_columns(columns)
208 @property
209 def column_type(self) -> ColumnType:
210 # Docstring inherited.
211 return self.operand.column_type
213 def __str__(self) -> str:
214 return f"{self.operand} DESC"
216 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
217 # Docstring inherited.
218 return visitor.visit_reversed(self)
221def validate_order_expression(expression: _ColumnExpression | Reversed) -> _ColumnExpression | Reversed:
222 """Check that a column expression can be used for sorting.
224 Parameters
225 ----------
226 expression : `OrderExpression`
227 Expression to check.
229 Returns
230 -------
231 expression : `OrderExpression`
232 The checked expression; returned to make this usable as a Pydantic
233 validator.
235 Raises
236 ------
237 InvalidQueryError
238 Raised if this expression is not one that can be used for sorting.
239 """
240 if expression.column_type not in ("int", "string", "float", "datetime"):
241 raise InvalidQueryError(f"Column type {expression.column_type} of {expression} is not ordered.")
242 return expression
245OrderExpression: TypeAlias = Annotated[
246 _ColumnExpression | Reversed,
247 pydantic.Field(discriminator="expression_type"),
248 pydantic.AfterValidator(validate_order_expression),
249]