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

1from collections.abc import Iterable, Mapping, Sequence 9 ctx1abcdefghi

2from typing import ( 9 ctx1abcdefghi

3 Any, 

4 Literal, 

5 Optional, 

6 TypeVar, 

7) 

8from typing import cast as typing_cast 9 ctx1abcdefghi

9 

10import sqlalchemy 9 ctx1abcdefghi

11from sqlalchemy import ( 9 ctx1abcdefghi

12 Column, 

13 ColumnElement, 

14 Extract, 

15 FunctionElement, 

16 FunctionFilter, 

17 Label, 

18 Over, 

19 TypeCoerce, 

20 WithinGroup, 

21) 

22from sqlalchemy.orm import InstrumentedAttribute, Mapped 9 ctx1abcdefghi

23from sqlalchemy.sql._typing import ( 9 ctx1abcdefghi

24 _ColumnExpressionArgument, 

25 _ColumnExpressionOrLiteralArgument, 

26 _ColumnExpressionOrStrLabelArgument, 

27) 

28from sqlalchemy.sql.elements import ( 9 ctx1abcdefghi

29 BinaryExpression, 

30 Case, 

31 Cast, 

32 CollectionAggregate, 

33 ColumnClause, 

34 TryCast, 

35 UnaryExpression, 

36) 

37from sqlalchemy.sql.type_api import TypeEngine 9 ctx1abcdefghi

38 

39from ._expression_select_cls import Select as Select 9 ctx1abcdefghi

40from ._expression_select_cls import SelectOfScalar as SelectOfScalar 9 ctx1abcdefghi

41from ._expression_select_gen import select as select 9 ctx1abcdefghi

42 

43_T = TypeVar("_T") 9 ctx1abcdefghi

44 

45_TypeEngineArgument = type[TypeEngine[_T]] | TypeEngine[_T] 9 ctx1abcdefghi

46 

47# Redefine operators that would only take a column expression to also take the (virtual) 

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

49 

50 

51def all_(expr: _ColumnExpressionArgument[_T] | _T) -> CollectionAggregate[bool]: 9 ctx1abcdefghi

52 return sqlalchemy.all_(expr) # ty: ignore[invalid-argument-type] 

53 

54 

55def and_( 9 ctx1abcdefghi

56 initial_clause: Literal[True] | _ColumnExpressionArgument[bool] | bool, 

57 *clauses: _ColumnExpressionArgument[bool] | bool, 

58) -> ColumnElement[bool]: 

59 return sqlalchemy.and_(initial_clause, *clauses) # ty: ignore[invalid-argument-type] 

60 

61 

62def any_(expr: _ColumnExpressionArgument[_T] | _T) -> CollectionAggregate[bool]: 9 ctx1abcdefghi

63 return sqlalchemy.any_(expr) # ty: ignore[invalid-argument-type] 

64 

65 

66def asc( 9 ctx1abcdefghi

67 column: _ColumnExpressionOrStrLabelArgument[_T] | _T, 

68) -> UnaryExpression[_T]: 

69 return sqlalchemy.asc(typing_cast(ColumnElement[_T], column)) 9 ctx1jklmnopqr

70 

71 

72def collate( 9 ctx1abcdefghi

73 expression: _ColumnExpressionArgument[str] | str, collation: str 

74) -> BinaryExpression[str]: 

75 return sqlalchemy.collate(expression, collation) # ty: ignore[invalid-argument-type] 

76 

77 

78def between( 9 ctx1abcdefghi

79 expr: _ColumnExpressionOrLiteralArgument[_T] | _T, 

80 lower_bound: Any, 

81 upper_bound: Any, 

82 symmetric: bool = False, 

83) -> BinaryExpression[bool]: 

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

85 

86 

87def not_(clause: _ColumnExpressionArgument[_T] | _T) -> ColumnElement[_T]: 9 ctx1abcdefghi

88 return sqlalchemy.not_(clause) # ty: ignore[no-matching-overload] 

89 

90 

91def case( 9 ctx1abcdefghi

92 *whens: tuple[_ColumnExpressionArgument[bool] | bool, Any] | Mapping[Any, Any], 

93 value: Any | None = None, 

94 else_: Any | None = None, 

95) -> Case[Any]: 

96 return sqlalchemy.case(*whens, value=value, else_=else_) # ty: ignore[invalid-argument-type] 

97 

98 

99def cast( 9 ctx1abcdefghi

100 expression: _ColumnExpressionOrLiteralArgument[Any] | Any, 

101 type_: "_TypeEngineArgument[_T]", 

102) -> Cast[_T]: 

103 return sqlalchemy.cast(expression, type_) 

104 

105 

106def try_cast( 9 ctx1abcdefghi

107 expression: _ColumnExpressionOrLiteralArgument[Any] | Any, 

108 type_: "_TypeEngineArgument[_T]", 

109) -> TryCast[_T]: 

110 return sqlalchemy.try_cast(expression, type_) 

111 

112 

113def desc( 9 ctx1abcdefghi

114 column: _ColumnExpressionOrStrLabelArgument[_T] | _T, 

115) -> UnaryExpression[_T]: 

116 return sqlalchemy.desc(typing_cast(ColumnElement[_T], column)) 9 ctx1jklmnopqr

117 

118 

119def distinct(expr: _ColumnExpressionArgument[_T] | _T) -> UnaryExpression[_T]: 9 ctx1abcdefghi

120 return sqlalchemy.distinct(expr) # ty: ignore[invalid-argument-type] 

121 

122 

123def bitwise_not(expr: _ColumnExpressionArgument[_T] | _T) -> UnaryExpression[_T]: 9 ctx1abcdefghi

124 return sqlalchemy.bitwise_not(expr) # ty: ignore[invalid-argument-type] 

125 

126 

127def extract(field: str, expr: _ColumnExpressionArgument[Any] | Any) -> Extract: 9 ctx1abcdefghi

128 return sqlalchemy.extract(field, expr) 

129 

130 

131def funcfilter( 9 ctx1abcdefghi

132 func: FunctionElement[_T], *criterion: _ColumnExpressionArgument[bool] | bool 

133) -> FunctionFilter[_T]: 

134 return sqlalchemy.funcfilter(func, *criterion) # ty: ignore[invalid-argument-type] 

135 

136 

137def label( 9 ctx1abcdefghi

138 name: str, 

139 element: _ColumnExpressionArgument[_T] | _T, 

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

141) -> Label[_T]: 

142 return sqlalchemy.label(name, element, type_=type_) # ty: ignore[invalid-argument-type] 

143 

144 

145def nulls_first( 9 ctx1abcdefghi

146 column: _ColumnExpressionArgument[_T] | _T, 

147) -> UnaryExpression[_T]: 

148 return sqlalchemy.nulls_first(typing_cast(ColumnElement[_T], column)) 9 ctx1jklmnopqr

149 

150 

151def nulls_last(column: _ColumnExpressionArgument[_T] | _T) -> UnaryExpression[_T]: 9 ctx1abcdefghi

152 return sqlalchemy.nulls_last(typing_cast(ColumnElement[_T], column)) 9 ctx1jklmnopqr

153 

154 

155def or_( 9 ctx1abcdefghi

156 initial_clause: Literal[False] | _ColumnExpressionArgument[bool] | bool, 

157 *clauses: _ColumnExpressionArgument[bool] | bool, 

158) -> ColumnElement[bool]: 

159 return sqlalchemy.or_(initial_clause, *clauses) # ty: ignore[invalid-argument-type] 9 ctx123456789!

160 

161 

162def over( 9 ctx1abcdefghi

163 element: FunctionElement[_T], 

164 partition_by: Iterable[_ColumnExpressionArgument[Any] | Any] 

165 | _ColumnExpressionArgument[Any] 

166 | Any 

167 | None = None, 

168 order_by: Iterable[_ColumnExpressionArgument[Any] | Any] 

169 | _ColumnExpressionArgument[Any] 

170 | Any 

171 | None = None, 

172 range_: tuple[int | None, int | None] | None = None, 

173 rows: tuple[int | None, int | None] | None = None, 

174) -> Over[_T]: 

175 return sqlalchemy.over( 

176 element, partition_by=partition_by, order_by=order_by, range_=range_, rows=rows 

177 ) 

178 

179 

180def tuple_( 9 ctx1abcdefghi

181 *clauses: _ColumnExpressionArgument[Any] | Any, 

182 types: Sequence["_TypeEngineArgument[Any]"] | None = None, 

183) -> sqlalchemy.Tuple: 

184 return sqlalchemy.tuple_(*clauses, types=types) 

185 

186 

187def type_coerce( 9 ctx1abcdefghi

188 expression: _ColumnExpressionOrLiteralArgument[Any] | Any, 

189 type_: "_TypeEngineArgument[_T]", 

190) -> TypeCoerce[_T]: 

191 return sqlalchemy.type_coerce(expression, type_) 

192 

193 

194def within_group( 9 ctx1abcdefghi

195 element: FunctionElement[_T], *order_by: _ColumnExpressionArgument[Any] | Any 

196) -> WithinGroup[_T]: 

197 return sqlalchemy.within_group(element, *order_by) 

198 

199 

200def col(column_expression: _T) -> Mapped[_T]: 9 ctx1abcdefghi

201 if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)): 36 ctx1stuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01

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

203 return column_expression # type: ignore 36 ctx1stuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01