_weakrefset.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205
  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),), self.__getstate__()
  70. def add(self, item):
  71. if self._pending_removals:
  72. self._commit_removals()
  73. self.data.add(ref(item, self._remove))
  74. def clear(self):
  75. if self._pending_removals:
  76. self._commit_removals()
  77. self.data.clear()
  78. def copy(self):
  79. return self.__class__(self)
  80. def pop(self):
  81. if self._pending_removals:
  82. self._commit_removals()
  83. while True:
  84. try:
  85. itemref = self.data.pop()
  86. except KeyError:
  87. raise KeyError('pop from empty WeakSet') from None
  88. item = itemref()
  89. if item is not None:
  90. return item
  91. def remove(self, item):
  92. if self._pending_removals:
  93. self._commit_removals()
  94. self.data.remove(ref(item))
  95. def discard(self, item):
  96. if self._pending_removals:
  97. self._commit_removals()
  98. self.data.discard(ref(item))
  99. def update(self, other):
  100. if self._pending_removals:
  101. self._commit_removals()
  102. for element in other:
  103. self.add(element)
  104. def __ior__(self, other):
  105. self.update(other)
  106. return self
  107. def difference(self, other):
  108. newset = self.copy()
  109. newset.difference_update(other)
  110. return newset
  111. __sub__ = difference
  112. def difference_update(self, other):
  113. self.__isub__(other)
  114. def __isub__(self, other):
  115. if self._pending_removals:
  116. self._commit_removals()
  117. if self is other:
  118. self.data.clear()
  119. else:
  120. self.data.difference_update(ref(item) for item in other)
  121. return self
  122. def intersection(self, other):
  123. return self.__class__(item for item in other if item in self)
  124. __and__ = intersection
  125. def intersection_update(self, other):
  126. self.__iand__(other)
  127. def __iand__(self, other):
  128. if self._pending_removals:
  129. self._commit_removals()
  130. self.data.intersection_update(ref(item) for item in other)
  131. return self
  132. def issubset(self, other):
  133. return self.data.issubset(ref(item) for item in other)
  134. __le__ = issubset
  135. def __lt__(self, other):
  136. return self.data < set(map(ref, other))
  137. def issuperset(self, other):
  138. return self.data.issuperset(ref(item) for item in other)
  139. __ge__ = issuperset
  140. def __gt__(self, other):
  141. return self.data > set(map(ref, other))
  142. def __eq__(self, other):
  143. if not isinstance(other, self.__class__):
  144. return NotImplemented
  145. return self.data == set(map(ref, other))
  146. def symmetric_difference(self, other):
  147. newset = self.copy()
  148. newset.symmetric_difference_update(other)
  149. return newset
  150. __xor__ = symmetric_difference
  151. def symmetric_difference_update(self, other):
  152. self.__ixor__(other)
  153. def __ixor__(self, other):
  154. if self._pending_removals:
  155. self._commit_removals()
  156. if self is other:
  157. self.data.clear()
  158. else:
  159. self.data.symmetric_difference_update(ref(item, self._remove) for item in other)
  160. return self
  161. def union(self, other):
  162. return self.__class__(e for s in (self, other) for e in s)
  163. __or__ = union
  164. def isdisjoint(self, other):
  165. return len(self.intersection(other)) == 0
  166. def __repr__(self):
  167. return repr(self.data)
  168. __class_getitem__ = classmethod(GenericAlias)