Coverage for dynamodx / keys.py: 85%

1from abc import ABC, abstractmethod 

2from dataclasses import dataclass 

3from typing import Any, Self 

4 

5 

6class Key(dict, ABC): 

7 @abstractmethod 

8 def expr_attr_names(self) -> dict[str, str]: ... 

9 

10 @abstractmethod 

11 def expr_attr_values(self) -> dict[str, Any]: ... 

12 

13 

14class SortKey(str): 

15 """ 

16 SortKey encapsulates a sort key value and optionally stores metadata 

17 used for nested data extraction and output transformation. 

18 

19 Parameters 

20 ---------- 

21 path_spec : str, optional 

22 JMESPath expression used to project nested data from the item. 

23 

24 rename_key : str, optional 

25 If provided, renames the sort key in the output. 

26 

27 projection_expr : str, optional 

28 Projection expression associated with the key. 

29 

30 expr_attr_names : dict[str, str], optional 

31 Attribute name mapping used with projection expressions. 

32 """ 

33 

34 __slots__ = ( 

35 'sk', 

36 'path_spec', 

37 'rename_key', 

38 'projection_expr', 

39 'expr_attr_names', 

40 ) 

41 

42 sk: str | None 

43 path_spec: str | None 

44 rename_key: str | None 

45 projection_expr: str | None 

46 expr_attr_names: dict[str, str] | None 

47 

48 def __new__( 

49 cls, 

50 *, 

51 path_spec: str | None = None, 

52 rename_key: str | None = None, 

53 projection_expr: str | None = None, 

54 expr_attr_names: dict[str, str] | None = None, 

55 **kwargs: str, 

56 ) -> Self: 

57 if len(kwargs) != 1: 

58 raise TypeError( 

59 f'SortKey() takes exactly one keyword argument ({len(kwargs)} given)' 

60 ) 

61 

62 ((name_sk, value_sk),) = kwargs.items() 

63 obj = super().__new__(cls, value_sk) 

64 

65 obj.sk = name_sk 

66 obj.path_spec = path_spec 

67 obj.rename_key = rename_key 

68 obj.projection_expr = projection_expr 

69 obj.expr_attr_names = expr_attr_names 

70 

71 return obj 

72 

73 

74class PartitionKey(Key): 

75 """Represents a partition key""" 

76 

77 def __init__( 

78 self, 

79 *, 

80 table_name: str | None = None, 

81 **kwargs, 

82 ) -> None: 

83 if len(kwargs) != 1: 

84 raise TypeError( 

85 f'PartitionKey() takes exactly one key=value argument ' 

86 f'({len(kwargs)} given)' 

87 ) 

88 

89 ((name_pk, value_pk),) = kwargs.items() 

90 super().__init__(**{name_pk: value_pk}) 

91 

92 self.name_pk = name_pk 

93 self.table_name = table_name 

94 

95 def expr_attr_names(self) -> dict[str, str]: 

96 return {'#pk': self.name_pk} 

97 

98 def expr_attr_values(self) -> dict[str, Any]: 

99 return {':pk': self[self.name_pk]} 

100 

101 def __add__(self, other: SortKey) -> 'PrimaryKey': 

102 pk = self.name_pk 

103 sk = other.sk 

104 kwargs = { 

105 pk: self[pk], 

106 sk: other, 

107 'table_name': self.table_name, 

108 } 

109 return PrimaryKey(**kwargs) 

110 

111 

112class PrimaryKey(Key): 

113 """Represents a composite key (partition key and sort key)""" 

114 

115 def __init__( 

116 self, 

117 *, 

118 table_name: str | None = None, 

119 **kwargs, 

120 ) -> None: 

121 """ 

122 Initializes a composite key using partition and sort key. 

123 

124 Parameters 

125 ---------- 

126 pk : str 

127 The partition key. 

128 sk : str 

129 The sort key. 

130 table_name : str, optional 

131 """ 

132 if len(kwargs) != 2: 

133 raise TypeError( 

134 f'PrimaryKey() takes exactly two key=value arguments ' 

135 f'({len(kwargs)} given)' 

136 ) 

137 

138 ( 

139 (name_pk, value_pk), 

140 (name_sk, value_sk), 

141 ) = kwargs.items() 

142 super().__init__(**{name_pk: value_pk, name_sk: value_sk}) 

143 

144 self.name_pk = name_pk 

145 self.name_sk = name_sk 

146 self.table_name = table_name 

147 

148 @property 

149 def sk(self) -> SortKey | str: 

150 return self[self.name_sk] 

151 

152 def expr_attr_names(self) -> dict[str, str]: 

153 return { 

154 '#pk': self.name_pk, 

155 '#sk': self.name_sk, 

156 } 

157 

158 def expr_attr_values(self) -> dict[str, Any]: 

159 return { 

160 ':pk': self[self.name_pk], 

161 ':sk': self[self.name_sk], 

162 } 

163 

164 def __add__(self, other: Self | SortKey) -> 'PrimaryKeySet': 

165 if isinstance(other, PrimaryKey): 

166 return PrimaryKeySet((self, other)) 

167 

168 if isinstance(other, SortKey): 

169 if other.sk is None: 

170 raise ValueError('SortKey name is required') 

171 

172 pk, sk = self.name_pk, other.sk 

173 kwargs = {pk: self[pk], sk: other} 

174 return PrimaryKeySet((self, PrimaryKey(**kwargs))) 

175 

176 return NotImplementedError 

177 

178 def __radd__(self, other: Any): 

179 if isinstance(other, PrimaryKeySet): 

180 return other + self 

181 

182 return NotImplementedError 

183 

184 

185@dataclass(frozen=True) 

186class PrimaryKeySet: 

187 pairs: tuple[PrimaryKey, ...] = () 

188 

189 def __add__(self, other: PrimaryKey | SortKey) -> 'PrimaryKeySet': 

190 if not isinstance(other, (PrimaryKey, SortKey)): 

191 return NotImplementedError 

192 

193 if isinstance(other, PrimaryKey): 

194 return PrimaryKeySet(pairs=self.pairs + (other,)) 

195 

196 if not self.pairs: 

197 raise ValueError('Cannot add SortKey to empty PrimaryKeySet') 

198 

199 last_pair = self.pairs[-1] 

200 next_pair = PrimaryKey( 

201 **{ 

202 last_pair.name_pk: last_pair[last_pair.name_pk], 

203 last_pair.name_sk: other, 

204 }, 

205 table_name=last_pair.table_name, 

206 ) 

207 

208 if next_pair in self.pairs: 

209 return self 

210 

211 return PrimaryKeySet(pairs=self.pairs + (next_pair,))