_weakrefset.py 5.8 KB

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