Coverage for dynamodx / conditions.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): 

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'{name} = {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, other: int) -> 'IfNotExistsExpr': 

89 return IfNotExistsExpr( 

90 path=self.path, 

91 value=self.value, 

92 r_value=other, 

93 operand='+', 

94 ) 

95 

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

97 return IfNotExistsExpr( 

98 path=self.path, 

99 value=self.value, 

100 r_value=other, 

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 Set(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 

126 to the SET action and can only 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 if isinstance(v, FuncExpr): 

137 self.func = v 

138 self.path = k 

139 self.value = v.value 

140 else: 

141 self.func = None 

142 self.path = k 

143 self.value = v 

144 self.operand = operand 

145 

146 def expr(self) -> str: 

147 if self.func: 

148 return self.func.expr() 

149 

150 operand = self.operand 

151 name = self.name_placeholder 

152 value = self.value_placeholder 

153 

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

155 # Incrementing and decrementing numeric attributes 

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

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

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

159 

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

161 

162 def expr_attr_values(self) -> dict: 

163 if self.func: 

164 return self.func.expr_attr_values() 

165 return super().expr_attr_values() 

166 

167 

168class Add(Expr): 

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

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

171 

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

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

174 

175 self.path = k 

176 self.value = v 

177 

178 def expr(self) -> str: 

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

180 

181 

182class Remove(Expr): 

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

184 self.path = path 

185 self.value = _Unset() 

186 

187 def expr(self) -> str: 

188 return self.name_placeholder 

189 

190 def expr_attr_values(self) -> dict: 

191 return {} 

192 

193 

194class Delete(Expr): 

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

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

197 

198 if not isinstance(v, set): 

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

200 

201 self.path = k 

202 self.value = v 

203 

204 def expr(self) -> str: 

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

206 

207 

208class UpdateExpr(dict): 

209 def __init__(self, *args) -> None: 

210 super().__init__() 

211 exprs = [x for x in args if x.value is not None] 

212 self.update(self.__asdict(exprs)) 

213 

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

215 expr_attr_names = reduce( 

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

217 ) 

218 expr_attr_values = reduce( 

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

220 ) 

221 

222 sets = list(filter(lambda attr: isinstance(attr, Set), exprs)) 

223 adds = list(filter(lambda attr: isinstance(attr, Add), exprs)) 

224 removes = list(filter(lambda attr: isinstance(attr, Remove), exprs)) 

225 deletes = list(filter(lambda attr: isinstance(attr, Delete), exprs)) 

226 

227 expr_parts = [] 

228 if sets: 

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

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

231 

232 if adds: 

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

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

235 

236 if removes: 

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

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

239 

240 if deletes: 

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

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

243 

244 update_expr = ' '.join(expr_parts) 

245 

246 return { 

247 'update_expr': update_expr, 

248 'expr_attr_names': expr_attr_names, 

249 'expr_attr_values': expr_attr_values, 

250 }