1from collections.abc import Set, Hashable
2import sys
3from typing import TypeVar, Generic
4from pyrsistent._pmap import pmap
5
6T_co = TypeVar('T_co', covariant=True)
7
8
9class PSet(Generic[T_co]):
10 """
11 Persistent set implementation. Built on top of the persistent map. The set supports all operations
12 in the Set protocol and is Hashable.
13
14 Do not instantiate directly, instead use the factory functions :py:func:`s` or :py:func:`pset`
15 to create an instance.
16
17 Random access and insert is log32(n) where n is the size of the set.
18
19 Some examples:
20
21 >>> s = pset([1, 2, 3, 1])
22 >>> s2 = s.add(4)
23 >>> s3 = s2.remove(2)
24 >>> s
25 pset([1, 2, 3])
26 >>> s2
27 pset([1, 2, 3, 4])
28 >>> s3
29 pset([1, 3, 4])
30 """
31 __slots__ = ('_map', '__weakref__')
32
33 def __new__(cls, m):
34 self = super(PSet, cls).__new__(cls)
35 self._map = m
36 return self
37
38 def __contains__(self, element):
39 return element in self._map
40
41 def __iter__(self):
42 return iter(self._map)
43
44 def __len__(self):
45 return len(self._map)
46
47 def __repr__(self):
48 if not self:
49 return 'p' + str(set(self))
50
51 return 'pset([{0}])'.format(str(set(self))[1:-1])
52
53 def __str__(self):
54 return self.__repr__()
55
56 def __hash__(self):
57 return hash(self._map)
58
59 def __reduce__(self):
60 # Pickling support
61 return pset, (list(self),)
62
63 @classmethod
64 def _from_iterable(cls, it, pre_size=8):
65 return PSet(pmap(dict((k, True) for k in it), pre_size=pre_size))
66
67 def add(self, element):
68 """
69 Return a new PSet with element added
70
71 >>> s1 = s(1, 2)
72 >>> s1.add(3)
73 pset([1, 2, 3])
74 """
75 return self.evolver().add(element).persistent()
76
77 def update(self, iterable):
78 """
79 Return a new PSet with elements in iterable added
80
81 >>> s1 = s(1, 2)
82 >>> s1.update([3, 4, 4])
83 pset([1, 2, 3, 4])
84 """
85 e = self.evolver()
86 for element in iterable:
87 e.add(element)
88
89 return e.persistent()
90
91 def remove(self, element):
92 """
93 Return a new PSet with element removed. Raises KeyError if element is not present.
94
95 >>> s1 = s(1, 2)
96 >>> s1.remove(2)
97 pset([1])
98 """
99 if element in self._map:
100 return self.evolver().remove(element).persistent()
101
102 raise KeyError("Element '%s' not present in PSet" % repr(element))
103
104 def discard(self, element):
105 """
106 Return a new PSet with element removed. Returns itself if element is not present.
107 """
108 if element in self._map:
109 return self.evolver().remove(element).persistent()
110
111 return self
112
113 class _Evolver(object):
114 __slots__ = ('_original_pset', '_pmap_evolver')
115
116 def __init__(self, original_pset):
117 self._original_pset = original_pset
118 self._pmap_evolver = original_pset._map.evolver()
119
120 def add(self, element):
121 self._pmap_evolver[element] = True
122 return self
123
124 def remove(self, element):
125 del self._pmap_evolver[element]
126 return self
127
128 def is_dirty(self):
129 return self._pmap_evolver.is_dirty()
130
131 def persistent(self):
132 if not self.is_dirty():
133 return self._original_pset
134
135 return PSet(self._pmap_evolver.persistent())
136
137 def __len__(self):
138 return len(self._pmap_evolver)
139
140 def copy(self):
141 return self
142
143 def evolver(self):
144 """
145 Create a new evolver for this pset. For a discussion on evolvers in general see the
146 documentation for the pvector evolver.
147
148 Create the evolver and perform various mutating updates to it:
149
150 >>> s1 = s(1, 2, 3)
151 >>> e = s1.evolver()
152 >>> _ = e.add(4)
153 >>> len(e)
154 4
155 >>> _ = e.remove(1)
156
157 The underlying pset remains the same:
158
159 >>> s1
160 pset([1, 2, 3])
161
162 The changes are kept in the evolver. An updated pmap can be created using the
163 persistent() function on the evolver.
164
165 >>> s2 = e.persistent()
166 >>> s2
167 pset([2, 3, 4])
168
169 The new pset will share data with the original pset in the same way that would have
170 been done if only using operations on the pset.
171 """
172 return PSet._Evolver(self)
173
174 # All the operations and comparisons you would expect on a set.
175 #
176 # This is not very beautiful. If we avoid inheriting from PSet we can use the
177 # __slots__ concepts (which requires a new style class) and hopefully save some memory.
178 __le__ = Set.__le__
179 __lt__ = Set.__lt__
180 __gt__ = Set.__gt__
181 __ge__ = Set.__ge__
182 __eq__ = Set.__eq__
183 __ne__ = Set.__ne__
184
185 __and__ = Set.__and__
186 __or__ = Set.__or__
187 __sub__ = Set.__sub__
188 __xor__ = Set.__xor__
189
190 issubset = __le__
191 issuperset = __ge__
192 union = __or__
193 intersection = __and__
194 difference = __sub__
195 symmetric_difference = __xor__
196
197 isdisjoint = Set.isdisjoint
198
199Set.register(PSet)
200Hashable.register(PSet)
201
202_EMPTY_PSET = PSet(pmap())
203
204
205def pset(iterable=(), pre_size=8):
206 """
207 Creates a persistent set from iterable. Optionally takes a sizing parameter equivalent to that
208 used for :py:func:`pmap`.
209
210 >>> s1 = pset([1, 2, 3, 2])
211 >>> s1
212 pset([1, 2, 3])
213 """
214 if not iterable:
215 return _EMPTY_PSET
216
217 return PSet._from_iterable(iterable, pre_size=pre_size)
218
219
220def s(*elements):
221 """
222 Create a persistent set.
223
224 Takes an arbitrary number of arguments to insert into the new set.
225
226 >>> s1 = s(1, 2, 3, 2)
227 >>> s1
228 pset([1, 2, 3])
229 """
230 return pset(elements)