Lib/_weakrefset.py

   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
   5 from _weakref import ref
   6
   7 __all__ = ['WeakSet']
   8
   9
  10 class _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
  35 class 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