_weakrefset.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
  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. from _weakref import ref
  5. __all__ = ['WeakSet']
  6. class _IterationGuard:
  7. # This context manager registers itself in the current iterators of the
  8. # weak container, such as to delay all removals until the context manager
  9. # exits.
  10. # This technique should be relatively thread-safe (since sets are).
  11. def __init__(self, weakcontainer):
  12. # Don't create cycles
  13. self.weakcontainer = ref(weakcontainer)
  14. def __enter__(self):
  15. w = self.weakcontainer()
  16. if w is not None:
  17. w._iterating.add(self)
  18. return self
  19. def __exit__(self, e, t, b):
  20. w = self.weakcontainer()
  21. if w is not None:
  22. s = w._iterating
  23. s.remove(self)
  24. if not s:
  25. w._commit_removals()
  26. class WeakSet:
  27. def __init__(self, data=None):
  28. self.data = set()
  29. def _remove(item, selfref=ref(self)):
  30. self = selfref()
  31. if self is not None:
  32. if self._iterating:
  33. self._pending_removals.append(item)
  34. else:
  35. self.data.discard(item)
  36. self._remove = _remove
  37. # A list of keys to be removed
  38. self._pending_removals = []
  39. self._iterating = set()
  40. if data is not None:
  41. self.update(data)
  42. def _commit_removals(self):
  43. l = self._pending_removals
  44. discard = self.data.discard
  45. while l:
  46. discard(l.pop())
  47. def __iter__(self):
  48. with _IterationGuard(self):
  49. for itemref in self.data:
  50. item = itemref()
  51. if item is not None:
  52. # Caveat: the iterator will keep a strong reference to
  53. # `item` until it is resumed or closed.
  54. yield item
  55. def __len__(self):
  56. return len(self.data) - len(self._pending_removals)
  57. def __contains__(self, item):
  58. try:
  59. wr = ref(item)
  60. except TypeError:
  61. return False
  62. return wr in self.data
  63. def __reduce__(self):
  64. return (self.__class__, (list(self),),
  65. getattr(self, '__dict__', None))
  66. def add(self, item):
  67. if self._pending_removals:
  68. self._commit_removals()
  69. self.data.add(ref(item, self._remove))
  70. def clear(self):
  71. if self._pending_removals:
  72. self._commit_removals()
  73. self.data.clear()
  74. def copy(self):
  75. return self.__class__(self)
  76. def pop(self):
  77. if self._pending_removals:
  78. self._commit_removals()
  79. while True:
  80. try:
  81. itemref = self.data.pop()
  82. except KeyError:
  83. raise KeyError('pop from empty WeakSet')
  84. item = itemref()
  85. if item is not None:
  86. return item
  87. def remove(self, item):
  88. if self._pending_removals:
  89. self._commit_removals()
  90. self.data.remove(ref(item))
  91. def discard(self, item):
  92. if self._pending_removals:
  93. self._commit_removals()
  94. self.data.discard(ref(item))
  95. def update(self, other):
  96. if self._pending_removals:
  97. self._commit_removals()
  98. for element in other:
  99. self.add(element)
  100. def __ior__(self, other):
  101. self.update(other)
  102. return self
  103. def difference(self, other):
  104. newset = self.copy()
  105. newset.difference_update(other)
  106. return newset
  107. __sub__ = difference
  108. def difference_update(self, other):
  109. self.__isub__(other)
  110. def __isub__(self, other):
  111. if self._pending_removals:
  112. self._commit_removals()
  113. if self is other:
  114. self.data.clear()
  115. else:
  116. self.data.difference_update(ref(item) for item in other)
  117. return self
  118. def intersection(self, other):
  119. return self.__class__(item for item in other if item in self)
  120. __and__ = intersection
  121. def intersection_update(self, other):
  122. self.__iand__(other)
  123. def __iand__(self, other):
  124. if self._pending_removals:
  125. self._commit_removals()
  126. self.data.intersection_update(ref(item) for item in other)
  127. return self
  128. def issubset(self, other):
  129. return self.data.issubset(ref(item) for item in other)
  130. __le__ = issubset
  131. def __lt__(self, other):
  132. return self.data < set(ref(item) for item in other)
  133. def issuperset(self, other):
  134. return self.data.issuperset(ref(item) for item in other)
  135. __ge__ = issuperset
  136. def __gt__(self, other):
  137. return self.data > set(ref(item) for item in other)
  138. def __eq__(self, other):
  139. if not isinstance(other, self.__class__):
  140. return NotImplemented
  141. return self.data == set(ref(item) for item in other)
  142. def symmetric_difference(self, other):
  143. newset = self.copy()
  144. newset.symmetric_difference_update(other)
  145. return newset
  146. __xor__ = symmetric_difference
  147. def symmetric_difference_update(self, other):
  148. self.__ixor__(other)
  149. def __ixor__(self, other):
  150. if self._pending_removals:
  151. self._commit_removals()
  152. if self is other:
  153. self.data.clear()
  154. else:
  155. self.data.symmetric_difference_update(ref(item, self._remove) for item in other)
  156. return self
  157. def union(self, other):
  158. return self.__class__(e for s in (self, other) for e in s)
  159. __or__ = union
  160. def isdisjoint(self, other):
  161. return len(self.intersection(other)) == 0