1# Access WeakSet through the weakref module.
2# This code is separated-out because it is needed
3# by abc.py to load everything else at startup.
4
5from _weakref import ref
6
7__all__ = ['WeakSet']
8
9
10class _IterationGuard(object):
11    # This context manager registers itself in the current iterators of the
12    # weak container, such as to delay all removals until the context manager
13    # exits.
14    # This technique should be relatively thread-safe (since sets are).
15
16    def __init__(self, weakcontainer):
17        # Don't create cycles
18        self.weakcontainer = ref(weakcontainer)
19
20    def __enter__(self):
21        w = self.weakcontainer()
22        if w is not None:
23            w._iterating.add(self)
24        return self
25
26    def __exit__(self, e, t, b):
27        w = self.weakcontainer()
28        if w is not None:
29            s = w._iterating
30            s.remove(self)
31            if not s:
32                w._commit_removals()
33
34
35class WeakSet(object):
36    def __init__(self, data=None):
37        self.data = set()
38        def _remove(item, selfref=ref(self)):
39            self = selfref()
40            if self is not None:
41                if self._iterating:
42                    self._pending_removals.append(item)
43                else:
44                    self.data.discard(item)
45        self._remove = _remove
46        # A list of keys to be removed
47        self._pending_removals = []
48        self._iterating = set()
49        if data is not None:
50            self.update(data)
51
52    def _commit_removals(self):
53        l = self._pending_removals
54        discard = self.data.discard
55        while l:
56            discard(l.pop())
57
58    def __iter__(self):
59        with _IterationGuard(self):
60            for itemref in self.data:
61                item = itemref()
62                if item is not None:
63                    yield item
64
65    def __len__(self):
66        return sum(x() is not None for x in self.data)
67
68    def __contains__(self, item):
69        return ref(item) in self.data
70
71    def __reduce__(self):
72        return (self.__class__, (list(self),),
73                getattr(self, '__dict__', None))
74
75    __hash__ = None
76
77    def add(self, item):
78        if self._pending_removals:
79            self._commit_removals()
80        self.data.add(ref(item, self._remove))
81
82    def clear(self):
83        if self._pending_removals:
84            self._commit_removals()
85        self.data.clear()
86
87    def copy(self):
88        return self.__class__(self)
89
90    def pop(self):
91        if self._pending_removals:
92            self._commit_removals()
93        while True:
94            try:
95                itemref = self.data.pop()
96            except KeyError:
97                raise KeyError('pop from empty WeakSet')
98            item = itemref()
99            if item is not None:
100                return item
101
102    def remove(self, item):
103        if self._pending_removals:
104            self._commit_removals()
105        self.data.remove(ref(item))
106
107    def discard(self, item):
108        if self._pending_removals:
109            self._commit_removals()
110        self.data.discard(ref(item))
111
112    def update(self, other):
113        if self._pending_removals:
114            self._commit_removals()
115        if isinstance(other, self.__class__):
116            self.data.update(other.data)
117        else:
118            for element in other:
119                self.add(element)
120
121    def __ior__(self, other):
122        self.update(other)
123        return self
124
125    # Helper functions for simple delegating methods.
126    def _apply(self, other, method):
127        if not isinstance(other, self.__class__):
128            other = self.__class__(other)
129        newdata = method(other.data)
130        newset = self.__class__()
131        newset.data = newdata
132        return newset
133
134    def difference(self, other):
135        return self._apply(other, self.data.difference)
136    __sub__ = difference
137
138    def difference_update(self, other):
139        if self._pending_removals:
140            self._commit_removals()
141        if self is other:
142            self.data.clear()
143        else:
144            self.data.difference_update(ref(item) for item in other)
145    def __isub__(self, other):
146        if self._pending_removals:
147            self._commit_removals()
148        if self is other:
149            self.data.clear()
150        else:
151            self.data.difference_update(ref(item) for item in other)
152        return self
153
154    def intersection(self, other):
155        return self._apply(other, self.data.intersection)
156    __and__ = intersection
157
158    def intersection_update(self, other):
159        if self._pending_removals:
160            self._commit_removals()
161        self.data.intersection_update(ref(item) for item in other)
162    def __iand__(self, other):
163        if self._pending_removals:
164            self._commit_removals()
165        self.data.intersection_update(ref(item) for item in other)
166        return self
167
168    def issubset(self, other):
169        return self.data.issubset(ref(item) for item in other)
170    __lt__ = issubset
171
172    def __le__(self, other):
173        return self.data <= set(ref(item) for item in other)
174
175    def issuperset(self, other):
176        return self.data.issuperset(ref(item) for item in other)
177    __gt__ = issuperset
178
179    def __ge__(self, other):
180        return self.data >= set(ref(item) for item in other)
181
182    def __eq__(self, other):
183        if not isinstance(other, self.__class__):
184            return NotImplemented
185        return self.data == set(ref(item) for item in other)
186
187    def symmetric_difference(self, other):
188        return self._apply(other, self.data.symmetric_difference)
189    __xor__ = symmetric_difference
190
191    def symmetric_difference_update(self, other):
192        if self._pending_removals:
193            self._commit_removals()
194        if self is other:
195            self.data.clear()
196        else:
197            self.data.symmetric_difference_update(ref(item) for item in other)
198    def __ixor__(self, other):
199        if self._pending_removals:
200            self._commit_removals()
201        if self is other:
202            self.data.clear()
203        else:
204            self.data.symmetric_difference_update(ref(item) for item in other)
205        return self
206
207    def union(self, other):
208        return self._apply(other, self.data.union)
209    __or__ = union
210
211    def isdisjoint(self, other):
212        return len(self.intersection(other)) == 0
213