Coverage for python/lsst/daf/butler/queries/tree/_column_expression.py: 48%

110 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-30 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/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ( 

31 "ColumnExpression", 

32 "OrderExpression", 

33 "UnaryExpression", 

34 "BinaryExpression", 

35 "Reversed", 

36 "UnaryOperator", 

37 "BinaryOperator", 

38 "validate_order_expression", 

39) 

40 

41from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias, TypeVar, final 

42 

43import pydantic 

44 

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 

51 

52if TYPE_CHECKING: 

53 from ..visitors import ColumnExpressionVisitor 

54 

55 

56_T = TypeVar("_T") 

57 

58 

59UnaryOperator: TypeAlias = Literal["-", "begin_of", "end_of"] 

60BinaryOperator: TypeAlias = Literal["+", "-", "*", "/", "%"] 

61 

62 

63@final 

64class UnaryExpression(ColumnExpressionBase): 

65 """A unary operation on a column expression that returns a non-boolean 

66 value. 

67 """ 

68 

69 expression_type: Literal["unary"] = "unary" 

70 

71 operand: ColumnExpression 

72 """Expression this one operates on.""" 

73 

74 operator: UnaryOperator 

75 """Operator this expression applies.""" 

76 

77 def gather_required_columns(self, columns: ColumnSet) -> None: 

78 # Docstring inherited. 

79 self.operand.gather_required_columns(columns) 

80 

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}.") 

90 

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" 

102 

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 

115 

116 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T: 

117 # Docstring inherited. 

118 return visitor.visit_unary_expression(self) 

119 

120 

121@final 

122class BinaryExpression(ColumnExpressionBase): 

123 """A binary operation on column expressions that returns a non-boolean 

124 value. 

125 """ 

126 

127 expression_type: Literal["binary"] = "binary" 

128 

129 a: ColumnExpression 

130 """Left-hand side expression this one operates on.""" 

131 

132 b: ColumnExpression 

133 """Right-hand side expression this one operates on.""" 

134 

135 operator: BinaryOperator 

136 """Operator this expression applies. 

137 

138 Integer '/' and '%' are defined as in SQL, not Python (though the 

139 definitions are the same for positive arguments). 

140 """ 

141 

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) 

146 

147 @property 

148 def column_type(self) -> ColumnType: 

149 # Docstring inherited. 

150 return self.a.column_type 

151 

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}" 

160 

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 

178 

179 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T: 

180 # Docstring inherited. 

181 return visitor.visit_binary_expression(self) 

182 

183 

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 

189 

190 

191ColumnExpression: TypeAlias = Annotated[_ColumnExpression, pydantic.Field(discriminator="expression_type")] 

192 

193 

194@final 

195class Reversed(ColumnExpressionBase): 

196 """A tag wrapper for `ColumnExpression` that indicates sorting in 

197 reverse order. 

198 """ 

199 

200 expression_type: Literal["reversed"] = "reversed" 

201 

202 operand: ColumnExpression 

203 """Expression to sort on in reverse.""" 

204 

205 def gather_required_columns(self, columns: ColumnSet) -> None: 

206 # Docstring inherited. 

207 self.operand.gather_required_columns(columns) 

208 

209 @property 

210 def column_type(self) -> ColumnType: 

211 # Docstring inherited. 

212 return self.operand.column_type 

213 

214 def __str__(self) -> str: 

215 return f"{self.operand} DESC" 

216 

217 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T: 

218 # Docstring inherited. 

219 return visitor.visit_reversed(self) 

220 

221 

222def validate_order_expression(expression: _ColumnExpression | Reversed) -> _ColumnExpression | Reversed: 

223 """Check that a column expression can be used for sorting. 

224 

225 Parameters 

226 ---------- 

227 expression : `OrderExpression` 

228 Expression to check. 

229 

230 Returns 

231 ------- 

232 expression : `OrderExpression` 

233 The checked expression; returned to make this usable as a Pydantic 

234 validator. 

235 

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 

244 

245 

246OrderExpression: TypeAlias = Annotated[ 

247 _ColumnExpression | Reversed, 

248 pydantic.Field(discriminator="expression_type"), 

249 pydantic.AfterValidator(validate_order_expression), 

250]