_weakrefset.py 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202
  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. l = self._pending_removals
  45. discard = self.data.discard
  46. while l:
  47. discard(l.pop())
  48. def __iter__(self):
  49. with _IterationGuard(self):
  50. for itemref in self.data:
  51. item = itemref()
  52. if item is not None:
  53. # Caveat: the iterator will keep a strong reference to
  54. # `item` until it is resumed or closed.
  55. yield item
  56. def __len__(self):
  57. return len(self.data) - len(self._pending_removals)
  58. def __contains__(self, item):
  59. try:
  60. wr = ref(item)
  61. except TypeError:
  62. return False
  63. return wr in self.data
  64. def __reduce__(self):
  65. return (self.__class__, (list(self),),
  66. getattr(self, '__dict__', None))
  67. def add(self, item):
  68. if self._pending_removals:
  69. self._commit_removals()
  70. self.data.add(ref(item, self._remove))
  71. def clear(self):
  72. if self._pending_removals:
  73. self._commit_removals()
  74. self.data.clear()
  75. def copy(self):
  76. return self.__class__(self)
  77. def pop(self):
  78. if self._pending_removals:
  79. self._commit_removals()
  80. while True:
  81. try:
  82. itemref = self.data.pop()
  83. except KeyError:
  84. raise KeyError('pop from empty WeakSet') from None
  85. item = itemref()
  86. if item is not None:
  87. return item
  88. def remove(self, item):
  89. if self._pending_removals:
  90. self._commit_removals()
  91. self.data.remove(ref(item))
  92. def discard(self, item):
  93. if self._pending_removals:
  94. self._commit_removals()
  95. self.data.discard(ref(item))
  96. def update(self, other):
  97. if self._pending_removals:
  98. self._commit_removals()
  99. for element in other:
  100. self.add(element)
  101. def __ior__(self, other):
  102. self.update(other)
  103. return self
  104. def difference(self, other):
  105. newset = self.copy()
  106. newset.difference_update(other)
  107. return newset
  108. __sub__ = difference
  109. def difference_update(self, other):
  110. self.__isub__(other)
  111. def __isub__(self, other):
  112. if self._pending_removals:
  113. self._commit_removals()
  114. if self is other:
  115. self.data.clear()
  116. else:
  117. self.data.difference_update(ref(item) for item in other)
  118. return self
  119. def intersection(self, other):
  120. return self.__class__(item for item in other if item in self)
  121. __and__ = intersection
  122. def intersection_update(self, other):
  123. self.__iand__(other)
  124. def __iand__(self, other):
  125. if self._pending_removals:
  126. self._commit_removals()
  127. self.data.intersection_update(ref(item) for item in other)
  128. return self
  129. def issubset(self, other):
  130. return self.data.issubset(ref(item) for item in other)
  131. __le__ = issubset
  132. def __lt__(self, other):
  133. return self.data < set(map(ref, other))
  134. def issuperset(self, other):
  135. return self.data.issuperset(ref(item) for item in other)
  136. __ge__ = issuperset
  137. def __gt__(self, other):
  138. return self.data > set(map(ref, other))
  139. def __eq__(self, other):
  140. if not isinstance(other, self.__class__):
  141. return NotImplemented
  142. return self.data == set(map(ref, other))
  143. def symmetric_difference(self, other):
  144. newset = self.copy()
  145. newset.symmetric_difference_update(other)
  146. return newset
  147. __xor__ = symmetric_difference
  148. def symmetric_difference_update(self, other):
  149. self.__ixor__(other)
  150. def __ixor__(self, other):
  151. if self._pending_removals:
  152. self._commit_removals()
  153. if self is other:
  154. self.data.clear()
  155. else:
  156. self.data.symmetric_difference_update(ref(item, self._remove) for item in other)
  157. return self
  158. def union(self, other):
  159. return self.__class__(e for s in (self, other) for e in s)
  160. __or__ = union
  161. def isdisjoint(self, other):
  162. return len(self.intersection(other)) == 0
  163. def __repr__(self):
  164. return repr(self.data)
  165. __class_getitem__ = classmethod(GenericAlias)