Coverage for sqlmodel/sql/expression.py: 63%

1from typing import ( 6 ctx1fabcde

2 Any, 

3 Iterable, 

4 Mapping, 

5 Optional, 

6 Sequence, 

7 Tuple, 

8 Type, 

9 TypeVar, 

10 Union, 

11) 

12 

13import sqlalchemy 6 ctx1fabcde

14from sqlalchemy import ( 6 ctx1fabcde

15 Column, 

16 ColumnElement, 

17 Extract, 

18 FunctionElement, 

19 FunctionFilter, 

20 Label, 

21 Over, 

22 TypeCoerce, 

23 WithinGroup, 

24) 

25from sqlalchemy.orm import InstrumentedAttribute, Mapped 6 ctx1fabcde

26from sqlalchemy.sql._typing import ( 6 ctx1fabcde

27 _ColumnExpressionArgument, 

28 _ColumnExpressionOrLiteralArgument, 

29 _ColumnExpressionOrStrLabelArgument, 

30) 

31from sqlalchemy.sql.elements import ( 6 ctx1fabcde

32 BinaryExpression, 

33 Case, 

34 Cast, 

35 CollectionAggregate, 

36 ColumnClause, 

37 TryCast, 

38 UnaryExpression, 

39) 

40from sqlalchemy.sql.type_api import TypeEngine 6 ctx1fabcde

41from typing_extensions import Literal 6 ctx1fabcde

42 

43from ._expression_select_cls import Select as Select 6 ctx1fabcde

44from ._expression_select_cls import SelectOfScalar as SelectOfScalar 6 ctx1fabcde

45from ._expression_select_gen import select as select 6 ctx1fabcde

46 

47_T = TypeVar("_T") 6 ctx1fabcde

48 

49_TypeEngineArgument = Union[Type[TypeEngine[_T]], TypeEngine[_T]] 6 ctx1fabcde

50 

51# Redefine operatos that would only take a column expresion to also take the (virtual) 

52# types of Pydantic models, e.g. str instead of only Mapped[str]. 

53 

54 

55def all_(expr: Union[_ColumnExpressionArgument[_T], _T]) -> CollectionAggregate[bool]: 6 ctx1fabcde

56 return sqlalchemy.all_(expr) # type: ignore[arg-type] 

57 

58 

59def and_( 5 ctx1abcde

60 initial_clause: Union[Literal[True], _ColumnExpressionArgument[bool], bool], 

61 *clauses: Union[_ColumnExpressionArgument[bool], bool], 

62) -> ColumnElement[bool]: 

63 return sqlalchemy.and_(initial_clause, *clauses) # type: ignore[arg-type] 

64 

65 

66def any_(expr: Union[_ColumnExpressionArgument[_T], _T]) -> CollectionAggregate[bool]: 6 ctx1fabcde

67 return sqlalchemy.any_(expr) # type: ignore[arg-type] 

68 

69 

70def asc( 5 ctx1abcde

71 column: Union[_ColumnExpressionOrStrLabelArgument[_T], _T], 

72) -> UnaryExpression[_T]: 

73 return sqlalchemy.asc(column) # type: ignore[arg-type] 

74 

75 

76def collate( 5 ctx1abcde

77 expression: Union[_ColumnExpressionArgument[str], str], collation: str 

78) -> BinaryExpression[str]: 

79 return sqlalchemy.collate(expression, collation) # type: ignore[arg-type] 

80 

81 

82def between( 5 ctx1abcde

83 expr: Union[_ColumnExpressionOrLiteralArgument[_T], _T], 

84 lower_bound: Any, 

85 upper_bound: Any, 

86 symmetric: bool = False, 

87) -> BinaryExpression[bool]: 

88 return sqlalchemy.between(expr, lower_bound, upper_bound, symmetric=symmetric) 

89 

90 

91def not_(clause: Union[_ColumnExpressionArgument[_T], _T]) -> ColumnElement[_T]: 6 ctx1fabcde

92 return sqlalchemy.not_(clause) # type: ignore[arg-type] 

93 

94 

95def case( 5 ctx1abcde

96 *whens: Union[ 

97 Tuple[Union[_ColumnExpressionArgument[bool], bool], Any], Mapping[Any, Any] 

98 ], 

99 value: Optional[Any] = None, 

100 else_: Optional[Any] = None, 

101) -> Case[Any]: 

102 return sqlalchemy.case(*whens, value=value, else_=else_) # type: ignore[arg-type] 

103 

104 

105def cast( 5 ctx1abcde

106 expression: Union[_ColumnExpressionOrLiteralArgument[Any], Any], 

107 type_: "_TypeEngineArgument[_T]", 

108) -> Cast[_T]: 

109 return sqlalchemy.cast(expression, type_) 

110 

111 

112def try_cast( 5 ctx1abcde

113 expression: Union[_ColumnExpressionOrLiteralArgument[Any], Any], 

114 type_: "_TypeEngineArgument[_T]", 

115) -> TryCast[_T]: 

116 return sqlalchemy.try_cast(expression, type_) 

117 

118 

119def desc( 5 ctx1abcde

120 column: Union[_ColumnExpressionOrStrLabelArgument[_T], _T], 

121) -> UnaryExpression[_T]: 

122 return sqlalchemy.desc(column) # type: ignore[arg-type] 

123 

124 

125def distinct(expr: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]: 6 ctx1fabcde

126 return sqlalchemy.distinct(expr) # type: ignore[arg-type] 

127 

128 

129def bitwise_not(expr: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]: 6 ctx1fabcde

130 return sqlalchemy.bitwise_not(expr) # type: ignore[arg-type] 

131 

132 

133def extract(field: str, expr: Union[_ColumnExpressionArgument[Any], Any]) -> Extract: 6 ctx1fabcde

134 return sqlalchemy.extract(field, expr) 

135 

136 

137def funcfilter( 5 ctx1abcde

138 func: FunctionElement[_T], *criterion: Union[_ColumnExpressionArgument[bool], bool] 

139) -> FunctionFilter[_T]: 

140 return sqlalchemy.funcfilter(func, *criterion) # type: ignore[arg-type] 

141 

142 

143def label( 5 ctx1abcde

144 name: str, 

145 element: Union[_ColumnExpressionArgument[_T], _T], 

146 type_: Optional["_TypeEngineArgument[_T]"] = None, 

147) -> Label[_T]: 

148 return sqlalchemy.label(name, element, type_=type_) # type: ignore[arg-type] 

149 

150 

151def nulls_first( 5 ctx1abcde

152 column: Union[_ColumnExpressionArgument[_T], _T], 

153) -> UnaryExpression[_T]: 

154 return sqlalchemy.nulls_first(column) # type: ignore[arg-type] 

155 

156 

157def nulls_last(column: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]: 6 ctx1fabcde

158 return sqlalchemy.nulls_last(column) # type: ignore[arg-type] 

159 

160 

161def or_( 5 ctx1abcde

162 initial_clause: Union[Literal[False], _ColumnExpressionArgument[bool], bool], 

163 *clauses: Union[_ColumnExpressionArgument[bool], bool], 

164) -> ColumnElement[bool]: 

165 return sqlalchemy.or_(initial_clause, *clauses) # type: ignore[arg-type] 9 ctx1pqrstuvwx

166 

167 

168def over( 5 ctx1abcde

169 element: FunctionElement[_T], 

170 partition_by: Optional[ 

171 Union[ 

172 Iterable[Union[_ColumnExpressionArgument[Any], Any]], 

173 _ColumnExpressionArgument[Any], 

174 Any, 

175 ] 

176 ] = None, 

177 order_by: Optional[ 

178 Union[ 

179 Iterable[Union[_ColumnExpressionArgument[Any], Any]], 

180 _ColumnExpressionArgument[Any], 

181 Any, 

182 ] 

183 ] = None, 

184 range_: Optional[Tuple[Optional[int], Optional[int]]] = None, 

185 rows: Optional[Tuple[Optional[int], Optional[int]]] = None, 

186) -> Over[_T]: 

187 return sqlalchemy.over( 

188 element, partition_by=partition_by, order_by=order_by, range_=range_, rows=rows 

189 ) 

190 

191 

192def tuple_( 5 ctx1abcde

193 *clauses: Union[_ColumnExpressionArgument[Any], Any], 

194 types: Optional[Sequence["_TypeEngineArgument[Any]"]] = None, 

195) -> Tuple[Any, ...]: 

196 return sqlalchemy.tuple_(*clauses, types=types) # type: ignore[return-value] 

197 

198 

199def type_coerce( 5 ctx1abcde

200 expression: Union[_ColumnExpressionOrLiteralArgument[Any], Any], 

201 type_: "_TypeEngineArgument[_T]", 

202) -> TypeCoerce[_T]: 

203 return sqlalchemy.type_coerce(expression, type_) 

204 

205 

206def within_group( 5 ctx1abcde

207 element: FunctionElement[_T], *order_by: Union[_ColumnExpressionArgument[Any], Any] 

208) -> WithinGroup[_T]: 

209 return sqlalchemy.within_group(element, *order_by) 

210 

211 

212def col(column_expression: _T) -> Mapped[_T]: 6 ctx1fabcde

213 if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)): 9 ctx1ghijklmno

214 raise RuntimeError(f"Not a SQLAlchemy column: {column_expression}") 

215 return column_expression # type: ignore 9 ctx1ghijklmno