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

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

2from typing import ( 9 ctx1abcdefghi

3 Any, 

4 Optional, 

5 TypeVar, 

6 Union, 

7) 

8 

9import sqlalchemy 9 ctx1abcdefghi

10from sqlalchemy import ( 9 ctx1abcdefghi

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 9 ctx1abcdefghi

22from sqlalchemy.sql._typing import ( 9 ctx1abcdefghi

23 _ColumnExpressionArgument, 

24 _ColumnExpressionOrLiteralArgument, 

25 _ColumnExpressionOrStrLabelArgument, 

26) 

27from sqlalchemy.sql.elements import ( 9 ctx1abcdefghi

28 BinaryExpression, 

29 Case, 

30 Cast, 

31 CollectionAggregate, 

32 ColumnClause, 

33 TryCast, 

34 UnaryExpression, 

35) 

36from sqlalchemy.sql.type_api import TypeEngine 9 ctx1abcdefghi

37from typing_extensions import Literal 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 = Union[type[TypeEngine[_T]], TypeEngine[_T]] 9 ctx1abcdefghi

46 

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

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

49 

50 

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

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

53 

54 

55def and_( 9 ctx1abcdefghi

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

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

58) -> ColumnElement[bool]: 

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

60 

61 

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

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

64 

65 

66def asc( 9 ctx1abcdefghi

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

68) -> UnaryExpression[_T]: 

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

70 

71 

72def collate( 9 ctx1abcdefghi

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

74) -> BinaryExpression[str]: 

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

76 

77 

78def between( 9 ctx1abcdefghi

79 expr: Union[_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: Union[_ColumnExpressionArgument[_T], _T]) -> ColumnElement[_T]: 9 ctx1abcdefghi

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

89 

90 

91def case( 9 ctx1abcdefghi

92 *whens: Union[ 

93 tuple[Union[_ColumnExpressionArgument[bool], bool], Any], Mapping[Any, Any] 

94 ], 

95 value: Optional[Any] = None, 

96 else_: Optional[Any] = None, 

97) -> Case[Any]: 

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

99 

100 

101def cast( 9 ctx1abcdefghi

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

103 type_: "_TypeEngineArgument[_T]", 

104) -> Cast[_T]: 

105 return sqlalchemy.cast(expression, type_) 

106 

107 

108def try_cast( 9 ctx1abcdefghi

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

110 type_: "_TypeEngineArgument[_T]", 

111) -> TryCast[_T]: 

112 return sqlalchemy.try_cast(expression, type_) 

113 

114 

115def desc( 9 ctx1abcdefghi

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

117) -> UnaryExpression[_T]: 

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

119 

120 

121def distinct(expr: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]: 9 ctx1abcdefghi

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

123 

124 

125def bitwise_not(expr: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]: 9 ctx1abcdefghi

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

127 

128 

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

130 return sqlalchemy.extract(field, expr) 

131 

132 

133def funcfilter( 9 ctx1abcdefghi

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

135) -> FunctionFilter[_T]: 

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

137 

138 

139def label( 9 ctx1abcdefghi

140 name: str, 

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

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

143) -> Label[_T]: 

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

145 

146 

147def nulls_first( 9 ctx1abcdefghi

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

149) -> UnaryExpression[_T]: 

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

151 

152 

153def nulls_last(column: Union[_ColumnExpressionArgument[_T], _T]) -> UnaryExpression[_T]: 9 ctx1abcdefghi

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

155 

156 

157def or_( 9 ctx1abcdefghi

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

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

160) -> ColumnElement[bool]: 

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

162 

163 

164def over( 9 ctx1abcdefghi

165 element: FunctionElement[_T], 

166 partition_by: Optional[ 

167 Union[ 

168 Iterable[Union[_ColumnExpressionArgument[Any], Any]], 

169 _ColumnExpressionArgument[Any], 

170 Any, 

171 ] 

172 ] = None, 

173 order_by: Optional[ 

174 Union[ 

175 Iterable[Union[_ColumnExpressionArgument[Any], Any]], 

176 _ColumnExpressionArgument[Any], 

177 Any, 

178 ] 

179 ] = None, 

180 range_: Optional[tuple[Optional[int], Optional[int]]] = None, 

181 rows: Optional[tuple[Optional[int], Optional[int]]] = None, 

182) -> Over[_T]: 

183 return sqlalchemy.over( 

184 element, partition_by=partition_by, order_by=order_by, range_=range_, rows=rows 

185 ) 

186 

187 

188def tuple_( 9 ctx1abcdefghi

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

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

191) -> tuple[Any, ...]: 

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

193 

194 

195def type_coerce( 9 ctx1abcdefghi

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

197 type_: "_TypeEngineArgument[_T]", 

198) -> TypeCoerce[_T]: 

199 return sqlalchemy.type_coerce(expression, type_) 

200 

201 

202def within_group( 9 ctx1abcdefghi

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

204) -> WithinGroup[_T]: 

205 return sqlalchemy.within_group(element, *order_by) 

206 

207 

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

209 if not isinstance(column_expression, (ColumnClause, Column, InstrumentedAttribute)): 9 ctx1jklmnopqr

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

211 return column_expression # type: ignore 9 ctx1jklmnopqr