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

1from collections.abc import Iterable, Mapping, Sequence 8 ctx1abcdefgh

2from typing import ( 8 ctx1abcdefgh

3 Any, 

4 Literal, 

5 Optional, 

6 TypeVar, 

7) 

8 

9import sqlalchemy 8 ctx1abcdefgh

10from sqlalchemy import ( 8 ctx1abcdefgh

11 Column, 

12 ColumnElement, 

13 Extract, 

14 FunctionElement, 

15 FunctionFilter, 

16 Label, 

17 Over, 

18 TypeCoerce, 

19 WithinGroup, 

20) 

21from sqlalchemy.orm import InstrumentedAttribute, Mapped 8 ctx1abcdefgh

22from sqlalchemy.sql._typing import ( 8 ctx1abcdefgh

23 _ColumnExpressionArgument, 

24 _ColumnExpressionOrLiteralArgument, 

25 _ColumnExpressionOrStrLabelArgument, 

26) 

27from sqlalchemy.sql.elements import ( 8 ctx1abcdefgh

28 BinaryExpression, 

29 Case, 

30 Cast, 

31 CollectionAggregate, 

32 ColumnClause, 

33 TryCast, 

34 UnaryExpression, 

35) 

36from sqlalchemy.sql.type_api import TypeEngine 8 ctx1abcdefgh

37 

38from ._expression_select_cls import Select as Select 8 ctx1abcdefgh

39from ._expression_select_cls import SelectOfScalar as SelectOfScalar 8 ctx1abcdefgh

40from ._expression_select_gen import select as select 8 ctx1abcdefgh

41 

42_T = TypeVar("_T") 8 ctx1abcdefgh

43 

44_TypeEngineArgument = type[TypeEngine[_T]] | TypeEngine[_T] 8 ctx1abcdefgh

45 

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

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

48 

49 

50def all_(expr: _ColumnExpressionArgument[_T] | _T) -> CollectionAggregate[bool]: 8 ctx1abcdefgh

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

52 

53 

54def and_( 8 ctx1abcdefgh

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

56 *clauses: _ColumnExpressionArgument[bool] | bool, 

57) -> ColumnElement[bool]: 

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

59 

60 

61def any_(expr: _ColumnExpressionArgument[_T] | _T) -> CollectionAggregate[bool]: 8 ctx1abcdefgh

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

63 

64 

65def asc( 8 ctx1abcdefgh

66 column: _ColumnExpressionOrStrLabelArgument[_T] | _T, 

67) -> UnaryExpression[_T]: 

68 return sqlalchemy.asc(column) # ty: ignore[invalid-argument-type] 

69 

70 

71def collate( 8 ctx1abcdefgh

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

73) -> BinaryExpression[str]: 

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

75 

76 

77def between( 8 ctx1abcdefgh

78 expr: _ColumnExpressionOrLiteralArgument[_T] | _T, 

79 lower_bound: Any, 

80 upper_bound: Any, 

81 symmetric: bool = False, 

82) -> BinaryExpression[bool]: 

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

84 

85 

86def not_(clause: _ColumnExpressionArgument[_T] | _T) -> ColumnElement[_T]: 8 ctx1abcdefgh

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

88 

89 

90def case( 8 ctx1abcdefgh

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

92 value: Any | None = None, 

93 else_: Any | None = None, 

94) -> Case[Any]: 

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

96 

97 

98def cast( 8 ctx1abcdefgh

99 expression: _ColumnExpressionOrLiteralArgument[Any] | Any, 

100 type_: "_TypeEngineArgument[_T]", 

101) -> Cast[_T]: 

102 return sqlalchemy.cast(expression, type_) 

103 

104 

105def try_cast( 8 ctx1abcdefgh

106 expression: _ColumnExpressionOrLiteralArgument[Any] | Any, 

107 type_: "_TypeEngineArgument[_T]", 

108) -> TryCast[_T]: 

109 return sqlalchemy.try_cast(expression, type_) 

110 

111 

112def desc( 8 ctx1abcdefgh

113 column: _ColumnExpressionOrStrLabelArgument[_T] | _T, 

114) -> UnaryExpression[_T]: 

115 return sqlalchemy.desc(column) # ty: ignore[invalid-argument-type] 

116 

117 

118def distinct(expr: _ColumnExpressionArgument[_T] | _T) -> UnaryExpression[_T]: 8 ctx1abcdefgh

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

120 

121 

122def bitwise_not(expr: _ColumnExpressionArgument[_T] | _T) -> UnaryExpression[_T]: 8 ctx1abcdefgh

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

124 

125 

126def extract(field: str, expr: _ColumnExpressionArgument[Any] | Any) -> Extract: 8 ctx1abcdefgh

127 return sqlalchemy.extract(field, expr) 

128 

129 

130def funcfilter( 8 ctx1abcdefgh

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

132) -> FunctionFilter[_T]: 

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

134 

135 

136def label( 8 ctx1abcdefgh

137 name: str, 

138 element: _ColumnExpressionArgument[_T] | _T, 

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

140) -> Label[_T]: 

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

142 

143 

144def nulls_first( 8 ctx1abcdefgh

145 column: _ColumnExpressionArgument[_T] | _T, 

146) -> UnaryExpression[_T]: 

147 return sqlalchemy.nulls_first(column) # ty: ignore[invalid-argument-type] 

148 

149 

150def nulls_last(column: _ColumnExpressionArgument[_T] | _T) -> UnaryExpression[_T]: 8 ctx1abcdefgh

151 return sqlalchemy.nulls_last(column) # ty: ignore[invalid-argument-type] 

152 

153 

154def or_( 8 ctx1abcdefgh

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

156 *clauses: _ColumnExpressionArgument[bool] | bool, 

157) -> ColumnElement[bool]: 

158 return sqlalchemy.or_(initial_clause, *clauses) # ty: ignore[invalid-argument-type] 8 ctx1yzABCDEF

159 

160 

161def over( 8 ctx1abcdefgh

162 element: FunctionElement[_T], 

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

164 | _ColumnExpressionArgument[Any] 

165 | Any 

166 | None = None, 

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

168 | _ColumnExpressionArgument[Any] 

169 | Any 

170 | None = None, 

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

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

173) -> Over[_T]: 

174 return sqlalchemy.over( 

175 element, partition_by=partition_by, order_by=order_by, range_=range_, rows=rows 

176 ) 

177 

178 

179def tuple_( 8 ctx1abcdefgh

180 *clauses: _ColumnExpressionArgument[Any] | Any, 

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

182) -> sqlalchemy.Tuple: 

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

184 

185 

186def type_coerce( 8 ctx1abcdefgh

187 expression: _ColumnExpressionOrLiteralArgument[Any] | Any, 

188 type_: "_TypeEngineArgument[_T]", 

189) -> TypeCoerce[_T]: 

190 return sqlalchemy.type_coerce(expression, type_) 

191 

192 

193def within_group( 8 ctx1abcdefgh

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

195) -> WithinGroup[_T]: 

196 return sqlalchemy.within_group(element, *order_by) 

197 

198 

199def col(column_expression: _T) -> Mapped[_T]: 8 ctx1abcdefgh

200 if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)): 16 ctx1ijklmnopqrstuvwx

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

202 return column_expression # type: ignore 16 ctx1ijklmnopqrstuvwx