Coverage for dynamodx / keys.py: 84%

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 

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

65 

66 obj.sk = name_sk 

67 obj.path_spec = path_spec 

68 obj.rename_key = rename_key 

69 obj.projection_expr = projection_expr 

70 obj.expr_attr_names = expr_attr_names 

71 

72 return obj 

73 

74 

75class PartitionKey(Key): 

76 """Represents a partition key""" 

77 

78 def __init__( 

79 self, 

80 *, 

81 table_name: str | None = None, 

82 **kwargs, 

83 ) -> None: 

84 if len(kwargs) != 1: 

85 raise TypeError( 

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

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

88 ) 

89 

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

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

92 

93 self.name_pk = name_pk 

94 self.table_name = table_name 

95 

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

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

98 

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

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

101 

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

103 pk = self.name_pk 

104 sk = other.sk 

105 kwargs = { 

106 pk: self[pk], 

107 sk: other, 

108 'table_name': self.table_name, 

109 } 

110 return PrimaryKey(**kwargs) 

111 

112 

113class PrimaryKey(Key): 

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

115 

116 def __init__( 

117 self, 

118 *, 

119 table_name: str | None = None, 

120 **kwargs, 

121 ) -> None: 

122 """ 

123 Initializes a composite key using partition and sort key. 

124 

125 Parameters 

126 ---------- 

127 pk : str 

128 The partition key. 

129 sk : str 

130 The sort key. 

131 table_name : str, optional 

132 """ 

133 if len(kwargs) != 2: 

134 raise TypeError( 

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

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

137 ) 

138 

139 (name_pk, value_pk), (name_sk, value_sk) = kwargs.items() 

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

141 

142 self.name_pk = name_pk 

143 self.name_sk = name_sk 

144 self.table_name = table_name 

145 

146 @property 

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

148 return self[self.name_sk] 

149 

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

151 return { 

152 '#pk': self.name_pk, 

153 '#sk': self.name_sk, 

154 } 

155 

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

157 sk = self[self.name_sk] 

158 return { 

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

160 ':sk': str(sk) if isinstance(sk, SortKey) else sk, 

161 } 

162 

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

164 if isinstance(other, PrimaryKey): 

165 return PrimaryKeySet((self, other)) 

166 

167 if isinstance(other, SortKey): 

168 if other.sk is None: 

169 raise ValueError('SortKey name is required') 

170 

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

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

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

174 

175 return NotImplemented 

176 

177 def __radd__(self, other: Any): 

178 if isinstance(other, PrimaryKeySet): 

179 return other + self 

180 

181 return NotImplemented 

182 

183 

184@dataclass(frozen=True) 

185class PrimaryKeySet: 

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

187 

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

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

190 return NotImplemented 

191 

192 if isinstance(other, PrimaryKey): 

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

194 

195 if not self.pairs: 

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

197 

198 last_pair = self.pairs[-1] 

199 kwargs = { 

200 last_pair.name_pk: last_pair[last_pair.name_pk], 

201 last_pair.name_sk: other, 

202 'table_name': last_pair.table_name, 

203 } 

204 next_pair = PrimaryKey(**kwargs) 

205 

206 if next_pair in self.pairs: 

207 return self 

208 

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