Coverage for dynamodx / expressions.py: 98%

1""" 

2- https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.UpdateExpressions.html 

3""" 

4 

5from abc import ABC, abstractmethod 

6from decimal import Decimal 

7from functools import reduce 

8from typing import Any, Literal 

9 

10 

11class _Unset: 

12 pass 

13 

14 

15class Expr(ABC): 

16 path: str 

17 value: str | set | Decimal | _Unset 

18 

19 def expr_attr_names(self) -> dict: 

20 return {self.name_placeholder: self.path} 

21 

22 def expr_attr_values(self) -> dict: 

23 return {self.value_placeholder: self.value} 

24 

25 @property 

26 def name_placeholder(self) -> str: 

27 return f'#n_{self.path}'.replace('.', '_') 

28 

29 @property 

30 def value_placeholder(self) -> str: 

31 return f':v_{self.path}' 

32 

33 @abstractmethod 

34 def expr(self) -> str: ... 

35 

36 

37class FuncExpr(Expr, ABC): 

38 def __init__( 

39 self, 

40 func: str, 

41 path: str, 

42 value: Any, 

43 ): 

44 self.func = func 

45 self.path = path 

46 self.value = value 

47 

48 def expr(self) -> str: 

49 func = self.func 

50 name = self.name_placeholder 

51 value = self.value_placeholder 

52 return f'{func}({name}, {value})' 

53 

54 

55class IfNotExistsExpr(FuncExpr): 

56 def __init__( 

57 self, 

58 path: str, 

59 value: Any, 

60 *, 

61 r_value: int | None = None, 

62 operand: Literal['+', '-'] | None = None, 

63 ): 

64 super().__init__('if_not_exists', path, value) 

65 

66 self.operand = operand 

67 self.r_value = r_value 

68 

69 def expr(self) -> str: 

70 expr = super().expr() 

71 operand = self.operand 

72 

73 if not self.r_value: 

74 return expr 

75 

76 return f'{expr} {operand} {self.value_placeholder}_r' 

77 

78 def expr_attr_values(self) -> dict: 

79 attrs = super().expr_attr_values() 

80 

81 if not self.r_value: 

82 return attrs 

83 

84 return attrs | { 

85 f'{self.value_placeholder}_r': self.r_value, 

86 } 

87 

88 def __add__(self, right_op: int) -> 'IfNotExistsExpr': 

89 return IfNotExistsExpr( 

90 path=self.path, 

91 value=self.value, 

92 r_value=right_op, 

93 operand='+', 

94 ) 

95 

96 def __sub__(self, right_op: int) -> 'IfNotExistsExpr': 

97 return IfNotExistsExpr( 

98 path=self.path, 

99 value=self.value, 

100 r_value=right_op, 

101 operand='-', 

102 ) 

103 

104 

105def list_append(**kwargs): 

106 (k, v), *_ = kwargs.items() 

107 return FuncExpr('list_append', k, v) 

108 

109 

110def if_not_exists(**kwargs): 

111 (k, v), *_ = kwargs.items() 

112 return IfNotExistsExpr(k, v) 

113 

114 

115class SetExpr(Expr): 

116 """ 

117 Use the `SET` action in an update expression to add one or more attributes 

118 to an item. 

119 

120 If any of these attributes already exists, they are overwritten by the new values. 

121 

122 If you want to avoid overwriting an existing attribute, you can use `SET` 

123 with the `if_not_exists` function. 

124 

125 The `if_not_exists` function is specific to the SET action and can only 

126 be used in an update expression. 

127 """ 

128 

129 def __init__( 

130 self, 

131 *, 

132 operand: Literal['=', '+', '-'] | None = None, 

133 **kwargs, 

134 ): 

135 (k, v), *_ = kwargs.items() 

136 

137 self.path = k 

138 self.value = v 

139 self.operand = operand 

140 

141 def expr(self) -> str: 

142 name = self.name_placeholder 

143 value = self.value_placeholder 

144 operand = self.operand 

145 

146 if isinstance(self.value, FuncExpr): 

147 expr = self.value.expr() 

148 return f'{name} = {expr}' 

149 

150 if operand in ('+', '-'): 

151 # Incrementing and decrementing numeric attributes 

152 # You can add to or subtract from an existing numeric attribute. 

153 # To do this, use the + (plus) and - (minus) operators. 

154 return f'{name} = {name} {operand} {value}' 

155 

156 return f'{name} = {value}' 

157 

158 def expr_attr_names(self) -> dict: 

159 attrs = super().expr_attr_names() 

160 

161 if isinstance(self.value, FuncExpr): 

162 return attrs | self.value.expr_attr_names() 

163 

164 return attrs 

165 

166 def expr_attr_values(self) -> dict: 

167 if isinstance(self.value, FuncExpr): 

168 return self.value.expr_attr_values() 

169 

170 return super().expr_attr_values() 

171 

172 

173class AddExpr(Expr): 

174 def __init__(self, **kwargs) -> None: 

175 (k, v), *_ = kwargs.items() 

176 

177 if not isinstance(v, (set, Decimal)): 

178 raise ValueError('ADD action supports only number and set data types') 

179 

180 self.path = k 

181 self.value = v 

182 

183 def expr(self) -> str: 

184 return f'{self.name_placeholder} {self.value_placeholder}' 

185 

186 

187class RemoveExpr(Expr): 

188 def __init__(self, path: str) -> None: 

189 self.path = path 

190 self.value = _Unset() 

191 

192 def expr(self) -> str: 

193 return self.name_placeholder 

194 

195 def expr_attr_values(self) -> dict: 

196 return {} 

197 

198 

199class DeleteExpr(Expr): 

200 def __init__(self, **kwargs) -> None: 

201 (k, v), *_ = kwargs.items() 

202 

203 if not isinstance(v, set): 

204 raise ValueError('DELETE action supports only Set data types') 

205 

206 self.path = k 

207 self.value = v 

208 

209 def expr(self) -> str: 

210 return f'{self.name_placeholder} {self.value_placeholder}' 

211 

212 

213class UpdateExpression(dict): 

214 def __init__(self, *args, exclude_none: bool = False) -> None: 

215 super().__init__() 

216 exprs = [x for x in args if not exclude_none or x.value is not None] 

217 self.update(self.__asdict(exprs)) 

218 

219 def __asdict(self, exprs: list[Expr] | None = None) -> dict: 

220 exprs = exprs or [] 

221 expr_attr_names = reduce( 

222 lambda acc, attr: {**acc, **attr.expr_attr_names()}, exprs, {} 

223 ) 

224 expr_attr_values = reduce( 

225 lambda acc, attr: {**acc, **attr.expr_attr_values()}, exprs, {} 

226 ) 

227 

228 sets = list(filter(lambda attr: isinstance(attr, SetExpr), exprs)) 

229 adds = list(filter(lambda attr: isinstance(attr, AddExpr), exprs)) 

230 removes = list(filter(lambda attr: isinstance(attr, RemoveExpr), exprs)) 

231 deletes = list(filter(lambda attr: isinstance(attr, DeleteExpr), exprs)) 

232 

233 expr_parts = [] 

234 if sets: 

235 set_expr = ', '.join(attr.expr() for attr in sets) 

236 expr_parts.append(f'SET {set_expr}') 

237 

238 if adds: 

239 add_expr = ', '.join(attr.expr() for attr in adds) 

240 expr_parts.append(f'ADD {add_expr}') 

241 

242 if removes: 

243 remove_expr = ', '.join(attr.expr() for attr in removes) 

244 expr_parts.append(f'REMOVE {remove_expr}') 

245 

246 if deletes: 

247 delete_expr = ', '.join(attr.expr() for attr in deletes) 

248 expr_parts.append(f'DELETE {delete_expr}') 

249 

250 update_expr = ' '.join(expr_parts) 

251 

252 return { 

253 'update_expr': update_expr, 

254 'expr_attr_names': expr_attr_names, 

255 'expr_attr_values': expr_attr_values, 

256 }