Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/pyrsistent/_pbag.py: 36%

1from collections.abc import Container, Iterable, Sized, Hashable 

2from functools import reduce 

3from pyrsistent._pmap import pmap 

4 

5 

6def _add_to_counters(counters, element): 

7 return counters.set(element, counters.get(element, 0) + 1) 

8 

9 

10class PBag(object): 

11 """ 

12 A persistent bag/multiset type. 

13 

14 Requires elements to be hashable, and allows duplicates, but has no 

15 ordering. Bags are hashable. 

16 

17 Do not instantiate directly, instead use the factory functions :py:func:`b` 

18 or :py:func:`pbag` to create an instance. 

19 

20 Some examples: 

21 

22 >>> s = pbag([1, 2, 3, 1]) 

23 >>> s2 = s.add(4) 

24 >>> s3 = s2.remove(1) 

25 >>> s 

26 pbag([1, 1, 2, 3]) 

27 >>> s2 

28 pbag([1, 1, 2, 3, 4]) 

29 >>> s3 

30 pbag([1, 2, 3, 4]) 

31 """ 

32 

33 __slots__ = ('_counts', '__weakref__') 

34 

35 def __init__(self, counts): 

36 self._counts = counts 

37 

38 def add(self, element): 

39 """ 

40 Add an element to the bag. 

41 

42 >>> s = pbag([1]) 

43 >>> s2 = s.add(1) 

44 >>> s3 = s.add(2) 

45 >>> s2 

46 pbag([1, 1]) 

47 >>> s3 

48 pbag([1, 2]) 

49 """ 

50 return PBag(_add_to_counters(self._counts, element)) 

51 

52 def update(self, iterable): 

53 """ 

54 Update bag with all elements in iterable. 

55 

56 >>> s = pbag([1]) 

57 >>> s.update([1, 2]) 

58 pbag([1, 1, 2]) 

59 """ 

60 if iterable: 

61 return PBag(reduce(_add_to_counters, iterable, self._counts)) 

62 

63 return self 

64 

65 def remove(self, element): 

66 """ 

67 Remove an element from the bag. 

68 

69 >>> s = pbag([1, 1, 2]) 

70 >>> s2 = s.remove(1) 

71 >>> s3 = s.remove(2) 

72 >>> s2 

73 pbag([1, 2]) 

74 >>> s3 

75 pbag([1, 1]) 

76 """ 

77 if element not in self._counts: 

78 raise KeyError(element) 

79 elif self._counts[element] == 1: 

80 newc = self._counts.remove(element) 

81 else: 

82 newc = self._counts.set(element, self._counts[element] - 1) 

83 return PBag(newc) 

84 

85 def count(self, element): 

86 """ 

87 Return the number of times an element appears. 

88 

89 

90 >>> pbag([]).count('non-existent') 

91 0 

92 >>> pbag([1, 1, 2]).count(1) 

93 2 

94 """ 

95 return self._counts.get(element, 0) 

96 

97 def __len__(self): 

98 """ 

99 Return the length including duplicates. 

100 

101 >>> len(pbag([1, 1, 2])) 

102 3 

103 """ 

104 return sum(self._counts.itervalues()) 

105 

106 def __iter__(self): 

107 """ 

108 Return an iterator of all elements, including duplicates. 

109 

110 >>> list(pbag([1, 1, 2])) 

111 [1, 1, 2] 

112 >>> list(pbag([1, 2])) 

113 [1, 2] 

114 """ 

115 for elt, count in self._counts.iteritems(): 

116 for i in range(count): 

117 yield elt 

118 

119 def __contains__(self, elt): 

120 """ 

121 Check if an element is in the bag. 

122 

123 >>> 1 in pbag([1, 1, 2]) 

124 True 

125 >>> 0 in pbag([1, 2]) 

126 False 

127 """ 

128 return elt in self._counts 

129 

130 def __repr__(self): 

131 return "pbag({0})".format(list(self)) 

132 

133 def __eq__(self, other): 

134 """ 

135 Check if two bags are equivalent, honoring the number of duplicates, 

136 and ignoring insertion order. 

137 

138 >>> pbag([1, 1, 2]) == pbag([1, 2]) 

139 False 

140 >>> pbag([2, 1, 0]) == pbag([0, 1, 2]) 

141 True 

142 """ 

143 if type(other) is not PBag: 

144 raise TypeError("Can only compare PBag with PBags") 

145 return self._counts == other._counts 

146 

147 def __lt__(self, other): 

148 raise TypeError('PBags are not orderable') 

149 

150 __le__ = __lt__ 

151 __gt__ = __lt__ 

152 __ge__ = __lt__ 

153 

154 # Multiset-style operations similar to collections.Counter 

155 

156 def __add__(self, other): 

157 """ 

158 Combine elements from two PBags. 

159 

160 >>> pbag([1, 2, 2]) + pbag([2, 3, 3]) 

161 pbag([1, 2, 2, 2, 3, 3]) 

162 """ 

163 if not isinstance(other, PBag): 

164 return NotImplemented 

165 result = self._counts.evolver() 

166 for elem, other_count in other._counts.iteritems(): 

167 result[elem] = self.count(elem) + other_count 

168 return PBag(result.persistent()) 

169 

170 def __sub__(self, other): 

171 """ 

172 Remove elements from one PBag that are present in another. 

173 

174 >>> pbag([1, 2, 2, 2, 3]) - pbag([2, 3, 3, 4]) 

175 pbag([1, 2, 2]) 

176 """ 

177 if not isinstance(other, PBag): 

178 return NotImplemented 

179 result = self._counts.evolver() 

180 for elem, other_count in other._counts.iteritems(): 

181 newcount = self.count(elem) - other_count 

182 if newcount > 0: 

183 result[elem] = newcount 

184 elif elem in self: 

185 result.remove(elem) 

186 return PBag(result.persistent()) 

187 

188 def __or__(self, other): 

189 """ 

190 Union: Keep elements that are present in either of two PBags. 

191 

192 >>> pbag([1, 2, 2, 2]) | pbag([2, 3, 3]) 

193 pbag([1, 2, 2, 2, 3, 3]) 

194 """ 

195 if not isinstance(other, PBag): 

196 return NotImplemented 

197 result = self._counts.evolver() 

198 for elem, other_count in other._counts.iteritems(): 

199 count = self.count(elem) 

200 newcount = max(count, other_count) 

201 result[elem] = newcount 

202 return PBag(result.persistent()) 

203 

204 def __and__(self, other): 

205 """ 

206 Intersection: Only keep elements that are present in both PBags. 

207 

208 >>> pbag([1, 2, 2, 2]) & pbag([2, 3, 3]) 

209 pbag([2]) 

210 """ 

211 if not isinstance(other, PBag): 

212 return NotImplemented 

213 result = pmap().evolver() 

214 for elem, count in self._counts.iteritems(): 

215 newcount = min(count, other.count(elem)) 

216 if newcount > 0: 

217 result[elem] = newcount 

218 return PBag(result.persistent()) 

219 

220 def __hash__(self): 

221 """ 

222 Hash based on value of elements. 

223 

224 >>> m = pmap({pbag([1, 2]): "it's here!"}) 

225 >>> m[pbag([2, 1])] 

226 "it's here!" 

227 >>> pbag([1, 1, 2]) in m 

228 False 

229 """ 

230 return hash(self._counts) 

231 

232 

233Container.register(PBag) 

234Iterable.register(PBag) 

235Sized.register(PBag) 

236Hashable.register(PBag) 

237 

238 

239def b(*elements): 

240 """ 

241 Construct a persistent bag. 

242 

243 Takes an arbitrary number of arguments to insert into the new persistent 

244 bag. 

245 

246 >>> b(1, 2, 3, 2) 

247 pbag([1, 2, 2, 3]) 

248 """ 

249 return pbag(elements) 

250 

251 

252def pbag(elements): 

253 """ 

254 Convert an iterable to a persistent bag. 

255 

256 Takes an iterable with elements to insert. 

257 

258 >>> pbag([1, 2, 3, 2]) 

259 pbag([1, 2, 2, 3]) 

260 """ 

261 if not elements: 

262 return _EMPTY_PBAG 

263 return PBag(reduce(_add_to_counters, elements, pmap())) 

264 

265 

266_EMPTY_PBAG = PBag(pmap()) 

267