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-03-26 02:48 -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, 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 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}.") 

89 

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" 

101 

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 

114 

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

116 # Docstring inherited. 

117 return visitor.visit_unary_expression(self) 

118 

119 

120@final 

121class BinaryExpression(ColumnExpressionBase): 

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

123 value. 

124 """ 

125 

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

127 

128 a: ColumnExpression 

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

130 

131 b: ColumnExpression 

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

133 

134 operator: BinaryOperator 

135 """Operator this expression applies. 

136 

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

138 definitions are the same for positive arguments). 

139 """ 

140 

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) 

145 

146 @property 

147 def column_type(self) -> ColumnType: 

148 # Docstring inherited. 

149 return self.a.column_type 

150 

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

159 

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 

177 

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

179 # Docstring inherited. 

180 return visitor.visit_binary_expression(self) 

181 

182 

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 = Union[ 

188 ColumnLiteral, 

189 ColumnReference, 

190 UnaryExpression, 

191 BinaryExpression, 

192] 

193 

194 

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

196 

197 

198@final 

199class Reversed(ColumnExpressionBase): 

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

201 reverse order. 

202 """ 

203 

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

205 

206 operand: ColumnExpression 

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

208 

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

210 # Docstring inherited. 

211 self.operand.gather_required_columns(columns) 

212 

213 @property 

214 def column_type(self) -> ColumnType: 

215 # Docstring inherited. 

216 return self.operand.column_type 

217 

218 def __str__(self) -> str: 

219 return f"{self.operand} DESC" 

220 

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

222 # Docstring inherited. 

223 return visitor.visit_reversed(self) 

224 

225 

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

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

228 

229 Parameters 

230 ---------- 

231 expression : `OrderExpression` 

232 Expression to check. 

233 

234 Returns 

235 ------- 

236 expression : `OrderExpression` 

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

238 validator. 

239 

240 Raises 

241 ------ 

242 InvalidQueryError 

243 Raised if this expression is not one that can be used for sorting. 

244 """ 

245 if expression.column_type not in ("int", "string", "float", "datetime"): 

246 raise InvalidQueryError(f"Column type {expression.column_type} of {expression} is not ordered.") 

247 return expression 

248 

249 

250OrderExpression: TypeAlias = Annotated[ 

251 Union[_ColumnExpression, Reversed], 

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

253 pydantic.AfterValidator(validate_order_expression), 

254]