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

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, Union, final 

42 

43import pydantic 

44 

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 

50 

51if TYPE_CHECKING: 

52 from ..visitors import ColumnExpressionVisitor 

53 

54 

55_T = TypeVar("_T") 

56 

57 

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

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

60 

61 

62@final 

63class UnaryExpression(ColumnExpressionBase): 

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

65 value. 

66 """ 

67 

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

69 

70 operand: ColumnExpression 

71 """Expression this one operates on.""" 

72 

73 operator: UnaryOperator 

74 """Operator this expression applies.""" 

75 

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

77 # Docstring inherited. 

78 self.operand.gather_required_columns(columns) 

79 

80 @property 

81 def precedence(self) -> int: 

82 # Docstring inherited. 

83 return 1 

84 

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

94 

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" 

104 

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 

117 

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

119 # Docstring inherited. 

120 return visitor.visit_unary_expression(self) 

121 

122 

123@final 

124class BinaryExpression(ColumnExpressionBase): 

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

126 value. 

127 """ 

128 

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

130 

131 a: ColumnExpression 

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

133 

134 b: ColumnExpression 

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

136 

137 operator: BinaryOperator 

138 """Operator this expression applies. 

139 

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

141 definitions are the same for positive arguments). 

142 """ 

143 

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) 

148 

149 @property 

150 def precedence(self) -> int: 

151 # Docstring inherited. 

152 match self.operator: 

153 case "*" | "/" | "%": 

154 return 2 

155 case "+" | "-": 

156 return 3 

157 

158 @property 

159 def column_type(self) -> ColumnType: 

160 # Docstring inherited. 

161 return self.a.column_type 

162 

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

178 

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 

196 

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

198 # Docstring inherited. 

199 return visitor.visit_binary_expression(self) 

200 

201 

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] 

212 

213 

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

215 

216 

217@final 

218class Reversed(ColumnExpressionBase): 

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

220 reverse order. 

221 """ 

222 

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

224 

225 operand: ColumnExpression 

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

227 

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

229 # Docstring inherited. 

230 self.operand.gather_required_columns(columns) 

231 

232 @property 

233 def precedence(self) -> int: 

234 # Docstring inherited. 

235 raise AssertionError("Order-reversed expressions can never be nested in other column expressions.") 

236 

237 @property 

238 def column_type(self) -> ColumnType: 

239 # Docstring inherited. 

240 return self.operand.column_type 

241 

242 def __str__(self) -> str: 

243 return f"{self.operand} DESC" 

244 

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

246 # Docstring inherited. 

247 return visitor.visit_reversed(self) 

248 

249 

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

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

252 

253 Parameters 

254 ---------- 

255 expression : `OrderExpression` 

256 Expression to check. 

257 

258 Returns 

259 ------- 

260 expression : `OrderExpression` 

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

262 validator. 

263 

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 

272 

273 

274OrderExpression: TypeAlias = Annotated[ 

275 Union[_ColumnExpression, Reversed], 

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

277 pydantic.AfterValidator(validate_order_expression), 

278]