weakref.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674
  1. """Weak reference support for Python.
  2. This module is an implementation of PEP 205:
  3. https://peps.python.org/pep-0205/
  4. """
  5. # Naming convention: Variables named "wr" are weak reference objects;
  6. # they are called this instead of "ref" to avoid name collisions with
  7. # the module-global ref() function imported from _weakref.
  8. from _weakref import (
  9. getweakrefcount,
  10. getweakrefs,
  11. ref,
  12. proxy,
  13. CallableProxyType,
  14. ProxyType,
  15. ReferenceType,
  16. _remove_dead_weakref)
  17. from _weakrefset import WeakSet, _IterationGuard
  18. import _collections_abc # Import after _weakref to avoid circular import.
  19. import sys
  20. import itertools
  21. ProxyTypes = (ProxyType, CallableProxyType)
  22. __all__ = ["ref", "proxy", "getweakrefcount", "getweakrefs",
  23. "WeakKeyDictionary", "ReferenceType", "ProxyType",
  24. "CallableProxyType", "ProxyTypes", "WeakValueDictionary",
  25. "WeakSet", "WeakMethod", "finalize"]
  26. _collections_abc.MutableSet.register(WeakSet)
  27. class WeakMethod(ref):
  28. """
  29. A custom `weakref.ref` subclass which simulates a weak reference to
  30. a bound method, working around the lifetime problem of bound methods.
  31. """
  32. __slots__ = "_func_ref", "_meth_type", "_alive", "__weakref__"
  33. def __new__(cls, meth, callback=None):
  34. try:
  35. obj = meth.__self__
  36. func = meth.__func__
  37. except AttributeError:
  38. raise TypeError("argument should be a bound method, not {}"
  39. .format(type(meth))) from None
  40. def _cb(arg):
  41. # The self-weakref trick is needed to avoid creating a reference
  42. # cycle.
  43. self = self_wr()
  44. if self._alive:
  45. self._alive = False
  46. if callback is not None:
  47. callback(self)
  48. self = ref.__new__(cls, obj, _cb)
  49. self._func_ref = ref(func, _cb)
  50. self._meth_type = type(meth)
  51. self._alive = True
  52. self_wr = ref(self)
  53. return self
  54. def __call__(self):
  55. obj = super().__call__()
  56. func = self._func_ref()
  57. if obj is None or func is None:
  58. return None
  59. return self._meth_type(func, obj)
  60. def __eq__(self, other):
  61. if isinstance(other, WeakMethod):
  62. if not self._alive or not other._alive:
  63. return self is other
  64. return ref.__eq__(self, other) and self._func_ref == other._func_ref
  65. return NotImplemented
  66. def __ne__(self, other):
  67. if isinstance(other, WeakMethod):
  68. if not self._alive or not other._alive:
  69. return self is not other
  70. return ref.__ne__(self, other) or self._func_ref != other._func_ref
  71. return NotImplemented
  72. __hash__ = ref.__hash__
  73. class WeakValueDictionary(_collections_abc.MutableMapping):
  74. """Mapping class that references values weakly.
  75. Entries in the dictionary will be discarded when no strong
  76. reference to the value exists anymore
  77. """
  78. # We inherit the constructor without worrying about the input
  79. # dictionary; since it uses our .update() method, we get the right
  80. # checks (if the other dictionary is a WeakValueDictionary,
  81. # objects are unwrapped on the way out, and we always wrap on the
  82. # way in).
  83. def __init__(self, other=(), /, **kw):
  84. def remove(wr, selfref=ref(self), _atomic_removal=_remove_dead_weakref):
  85. self = selfref()
  86. if self is not None:
  87. if self._iterating:
  88. self._pending_removals.append(wr.key)
  89. else:
  90. # Atomic removal is necessary since this function
  91. # can be called asynchronously by the GC
  92. _atomic_removal(self.data, wr.key)
  93. self._remove = remove
  94. # A list of keys to be removed
  95. self._pending_removals = []
  96. self._iterating = set()
  97. self.data = {}
  98. self.update(other, **kw)
  99. def _commit_removals(self, _atomic_removal=_remove_dead_weakref):
  100. pop = self._pending_removals.pop
  101. d = self.data
  102. # We shouldn't encounter any KeyError, because this method should
  103. # always be called *before* mutating the dict.
  104. while True:
  105. try:
  106. key = pop()
  107. except IndexError:
  108. return
  109. _atomic_removal(d, key)
  110. def __getitem__(self, key):
  111. if self._pending_removals:
  112. self._commit_removals()
  113. o = self.data[key]()
  114. if o is None:
  115. raise KeyError(key)
  116. else:
  117. return o
  118. def __delitem__(self, key):
  119. if self._pending_removals:
  120. self._commit_removals()
  121. del self.data[key]
  122. def __len__(self):
  123. if self._pending_removals:
  124. self._commit_removals()
  125. return len(self.data)
  126. def __contains__(self, key):
  127. if self._pending_removals:
  128. self._commit_removals()
  129. try:
  130. o = self.data[key]()
  131. except KeyError:
  132. return False
  133. return o is not None
  134. def __repr__(self):
  135. return "<%s at %#x>" % (self.__class__.__name__, id(self))
  136. def __setitem__(self, key, value):
  137. if self._pending_removals:
  138. self._commit_removals()
  139. self.data[key] = KeyedRef(value, self._remove, key)
  140. def copy(self):
  141. if self._pending_removals:
  142. self._commit_removals()
  143. new = WeakValueDictionary()
  144. with _IterationGuard(self):
  145. for key, wr in self.data.items():
  146. o = wr()
  147. if o is not None:
  148. new[key] = o
  149. return new
  150. __copy__ = copy
  151. def __deepcopy__(self, memo):
  152. from copy import deepcopy
  153. if self._pending_removals:
  154. self._commit_removals()
  155. new = self.__class__()
  156. with _IterationGuard(self):
  157. for key, wr in self.data.items():
  158. o = wr()
  159. if o is not None:
  160. new[deepcopy(key, memo)] = o
  161. return new
  162. def get(self, key, default=None):
  163. if self._pending_removals:
  164. self._commit_removals()
  165. try:
  166. wr = self.data[key]
  167. except KeyError:
  168. return default
  169. else:
  170. o = wr()
  171. if o is None:
  172. # This should only happen
  173. return default
  174. else:
  175. return o
  176. def items(self):
  177. if self._pending_removals:
  178. self._commit_removals()
  179. with _IterationGuard(self):
  180. for k, wr in self.data.items():
  181. v = wr()
  182. if v is not None:
  183. yield k, v
  184. def keys(self):
  185. if self._pending_removals:
  186. self._commit_removals()
  187. with _IterationGuard(self):
  188. for k, wr in self.data.items():
  189. if wr() is not None:
  190. yield k
  191. __iter__ = keys
  192. def itervaluerefs(self):
  193. """Return an iterator that yields the weak references to the values.
  194. The references are not guaranteed to be 'live' at the time
  195. they are used, so the result of calling the references needs
  196. to be checked before being used. This can be used to avoid
  197. creating references that will cause the garbage collector to
  198. keep the values around longer than needed.
  199. """
  200. if self._pending_removals:
  201. self._commit_removals()
  202. with _IterationGuard(self):
  203. yield from self.data.values()
  204. def values(self):
  205. if self._pending_removals:
  206. self._commit_removals()
  207. with _IterationGuard(self):
  208. for wr in self.data.values():
  209. obj = wr()
  210. if obj is not None:
  211. yield obj
  212. def popitem(self):
  213. if self._pending_removals:
  214. self._commit_removals()
  215. while True:
  216. key, wr = self.data.popitem()
  217. o = wr()
  218. if o is not None:
  219. return key, o
  220. def pop(self, key, *args):
  221. if self._pending_removals:
  222. self._commit_removals()
  223. try:
  224. o = self.data.pop(key)()
  225. except KeyError:
  226. o = None
  227. if o is None:
  228. if args:
  229. return args[0]
  230. else:
  231. raise KeyError(key)
  232. else:
  233. return o
  234. def setdefault(self, key, default=None):
  235. try:
  236. o = self.data[key]()
  237. except KeyError:
  238. o = None
  239. if o is None:
  240. if self._pending_removals:
  241. self._commit_removals()
  242. self.data[key] = KeyedRef(default, self._remove, key)
  243. return default
  244. else:
  245. return o
  246. def update(self, other=None, /, **kwargs):
  247. if self._pending_removals:
  248. self._commit_removals()
  249. d = self.data
  250. if other is not None:
  251. if not hasattr(other, "items"):
  252. other = dict(other)
  253. for key, o in other.items():
  254. d[key] = KeyedRef(o, self._remove, key)
  255. for key, o in kwargs.items():
  256. d[key] = KeyedRef(o, self._remove, key)
  257. def valuerefs(self):
  258. """Return a list of weak references to the values.
  259. The references are not guaranteed to be 'live' at the time
  260. they are used, so the result of calling the references needs
  261. to be checked before being used. This can be used to avoid
  262. creating references that will cause the garbage collector to
  263. keep the values around longer than needed.
  264. """
  265. if self._pending_removals:
  266. self._commit_removals()
  267. return list(self.data.values())
  268. def __ior__(self, other):
  269. self.update(other)
  270. return self
  271. def __or__(self, other):
  272. if isinstance(other, _collections_abc.Mapping):
  273. c = self.copy()
  274. c.update(other)
  275. return c
  276. return NotImplemented
  277. def __ror__(self, other):
  278. if isinstance(other, _collections_abc.Mapping):
  279. c = self.__class__()
  280. c.update(other)
  281. c.update(self)
  282. return c
  283. return NotImplemented
  284. class KeyedRef(ref):
  285. """Specialized reference that includes a key corresponding to the value.
  286. This is used in the WeakValueDictionary to avoid having to create
  287. a function object for each key stored in the mapping. A shared
  288. callback object can use the 'key' attribute of a KeyedRef instead
  289. of getting a reference to the key from an enclosing scope.
  290. """
  291. __slots__ = "key",
  292. def __new__(type, ob, callback, key):
  293. self = ref.__new__(type, ob, callback)
  294. self.key = key
  295. return self
  296. def __init__(self, ob, callback, key):
  297. super().__init__(ob, callback)
  298. class WeakKeyDictionary(_collections_abc.MutableMapping):
  299. """ Mapping class that references keys weakly.
  300. Entries in the dictionary will be discarded when there is no
  301. longer a strong reference to the key. This can be used to
  302. associate additional data with an object owned by other parts of
  303. an application without adding attributes to those objects. This
  304. can be especially useful with objects that override attribute
  305. accesses.
  306. """
  307. def __init__(self, dict=None):
  308. self.data = {}
  309. def remove(k, selfref=ref(self)):
  310. self = selfref()
  311. if self is not None:
  312. if self._iterating:
  313. self._pending_removals.append(k)
  314. else:
  315. try:
  316. del self.data[k]
  317. except KeyError:
  318. pass
  319. self._remove = remove
  320. # A list of dead weakrefs (keys to be removed)
  321. self._pending_removals = []
  322. self._iterating = set()
  323. self._dirty_len = False
  324. if dict is not None:
  325. self.update(dict)
  326. def _commit_removals(self):
  327. # NOTE: We don't need to call this method before mutating the dict,
  328. # because a dead weakref never compares equal to a live weakref,
  329. # even if they happened to refer to equal objects.
  330. # However, it means keys may already have been removed.
  331. pop = self._pending_removals.pop
  332. d = self.data
  333. while True:
  334. try:
  335. key = pop()
  336. except IndexError:
  337. return
  338. try:
  339. del d[key]
  340. except KeyError:
  341. pass
  342. def _scrub_removals(self):
  343. d = self.data
  344. self._pending_removals = [k for k in self._pending_removals if k in d]
  345. self._dirty_len = False
  346. def __delitem__(self, key):
  347. self._dirty_len = True
  348. del self.data[ref(key)]
  349. def __getitem__(self, key):
  350. return self.data[ref(key)]
  351. def __len__(self):
  352. if self._dirty_len and self._pending_removals:
  353. # self._pending_removals may still contain keys which were
  354. # explicitly removed, we have to scrub them (see issue #21173).
  355. self._scrub_removals()
  356. return len(self.data) - len(self._pending_removals)
  357. def __repr__(self):
  358. return "<%s at %#x>" % (self.__class__.__name__, id(self))
  359. def __setitem__(self, key, value):
  360. self.data[ref(key, self._remove)] = value
  361. def copy(self):
  362. new = WeakKeyDictionary()
  363. with _IterationGuard(self):
  364. for key, value in self.data.items():
  365. o = key()
  366. if o is not None:
  367. new[o] = value
  368. return new
  369. __copy__ = copy
  370. def __deepcopy__(self, memo):
  371. from copy import deepcopy
  372. new = self.__class__()
  373. with _IterationGuard(self):
  374. for key, value in self.data.items():
  375. o = key()
  376. if o is not None:
  377. new[o] = deepcopy(value, memo)
  378. return new
  379. def get(self, key, default=None):
  380. return self.data.get(ref(key),default)
  381. def __contains__(self, key):
  382. try:
  383. wr = ref(key)
  384. except TypeError:
  385. return False
  386. return wr in self.data
  387. def items(self):
  388. with _IterationGuard(self):
  389. for wr, value in self.data.items():
  390. key = wr()
  391. if key is not None:
  392. yield key, value
  393. def keys(self):
  394. with _IterationGuard(self):
  395. for wr in self.data:
  396. obj = wr()
  397. if obj is not None:
  398. yield obj
  399. __iter__ = keys
  400. def values(self):
  401. with _IterationGuard(self):
  402. for wr, value in self.data.items():
  403. if wr() is not None:
  404. yield value
  405. def keyrefs(self):
  406. """Return a list of weak references to the keys.
  407. The references are not guaranteed to be 'live' at the time
  408. they are used, so the result of calling the references needs
  409. to be checked before being used. This can be used to avoid
  410. creating references that will cause the garbage collector to
  411. keep the keys around longer than needed.
  412. """
  413. return list(self.data)
  414. def popitem(self):
  415. self._dirty_len = True
  416. while True:
  417. key, value = self.data.popitem()
  418. o = key()
  419. if o is not None:
  420. return o, value
  421. def pop(self, key, *args):
  422. self._dirty_len = True
  423. return self.data.pop(ref(key), *args)
  424. def setdefault(self, key, default=None):
  425. return self.data.setdefault(ref(key, self._remove),default)
  426. def update(self, dict=None, /, **kwargs):
  427. d = self.data
  428. if dict is not None:
  429. if not hasattr(dict, "items"):
  430. dict = type({})(dict)
  431. for key, value in dict.items():
  432. d[ref(key, self._remove)] = value
  433. if len(kwargs):
  434. self.update(kwargs)
  435. def __ior__(self, other):
  436. self.update(other)
  437. return self
  438. def __or__(self, other):
  439. if isinstance(other, _collections_abc.Mapping):
  440. c = self.copy()
  441. c.update(other)
  442. return c
  443. return NotImplemented
  444. def __ror__(self, other):
  445. if isinstance(other, _collections_abc.Mapping):
  446. c = self.__class__()
  447. c.update(other)
  448. c.update(self)
  449. return c
  450. return NotImplemented
  451. class finalize:
  452. """Class for finalization of weakrefable objects
  453. finalize(obj, func, *args, **kwargs) returns a callable finalizer
  454. object which will be called when obj is garbage collected. The
  455. first time the finalizer is called it evaluates func(*arg, **kwargs)
  456. and returns the result. After this the finalizer is dead, and
  457. calling it just returns None.
  458. When the program exits any remaining finalizers for which the
  459. atexit attribute is true will be run in reverse order of creation.
  460. By default atexit is true.
  461. """
  462. # Finalizer objects don't have any state of their own. They are
  463. # just used as keys to lookup _Info objects in the registry. This
  464. # ensures that they cannot be part of a ref-cycle.
  465. __slots__ = ()
  466. _registry = {}
  467. _shutdown = False
  468. _index_iter = itertools.count()
  469. _dirty = False
  470. _registered_with_atexit = False
  471. class _Info:
  472. __slots__ = ("weakref", "func", "args", "kwargs", "atexit", "index")
  473. def __init__(self, obj, func, /, *args, **kwargs):
  474. if not self._registered_with_atexit:
  475. # We may register the exit function more than once because
  476. # of a thread race, but that is harmless
  477. import atexit
  478. atexit.register(self._exitfunc)
  479. finalize._registered_with_atexit = True
  480. info = self._Info()
  481. info.weakref = ref(obj, self)
  482. info.func = func
  483. info.args = args
  484. info.kwargs = kwargs or None
  485. info.atexit = True
  486. info.index = next(self._index_iter)
  487. self._registry[self] = info
  488. finalize._dirty = True
  489. def __call__(self, _=None):
  490. """If alive then mark as dead and return func(*args, **kwargs);
  491. otherwise return None"""
  492. info = self._registry.pop(self, None)
  493. if info and not self._shutdown:
  494. return info.func(*info.args, **(info.kwargs or {}))
  495. def detach(self):
  496. """If alive then mark as dead and return (obj, func, args, kwargs);
  497. otherwise return None"""
  498. info = self._registry.get(self)
  499. obj = info and info.weakref()
  500. if obj is not None and self._registry.pop(self, None):
  501. return (obj, info.func, info.args, info.kwargs or {})
  502. def peek(self):
  503. """If alive then return (obj, func, args, kwargs);
  504. otherwise return None"""
  505. info = self._registry.get(self)
  506. obj = info and info.weakref()
  507. if obj is not None:
  508. return (obj, info.func, info.args, info.kwargs or {})
  509. @property
  510. def alive(self):
  511. """Whether finalizer is alive"""
  512. return self in self._registry
  513. @property
  514. def atexit(self):
  515. """Whether finalizer should be called at exit"""
  516. info = self._registry.get(self)
  517. return bool(info) and info.atexit
  518. @atexit.setter
  519. def atexit(self, value):
  520. info = self._registry.get(self)
  521. if info:
  522. info.atexit = bool(value)
  523. def __repr__(self):
  524. info = self._registry.get(self)
  525. obj = info and info.weakref()
  526. if obj is None:
  527. return '<%s object at %#x; dead>' % (type(self).__name__, id(self))
  528. else:
  529. return '<%s object at %#x; for %r at %#x>' % \
  530. (type(self).__name__, id(self), type(obj).__name__, id(obj))
  531. @classmethod
  532. def _select_for_exit(cls):
  533. # Return live finalizers marked for exit, oldest first
  534. L = [(f,i) for (f,i) in cls._registry.items() if i.atexit]
  535. L.sort(key=lambda item:item[1].index)
  536. return [f for (f,i) in L]
  537. @classmethod
  538. def _exitfunc(cls):
  539. # At shutdown invoke finalizers for which atexit is true.
  540. # This is called once all other non-daemonic threads have been
  541. # joined.
  542. reenable_gc = False
  543. try:
  544. if cls._registry:
  545. import gc
  546. if gc.isenabled():
  547. reenable_gc = True
  548. gc.disable()
  549. pending = None
  550. while True:
  551. if pending is None or finalize._dirty:
  552. pending = cls._select_for_exit()
  553. finalize._dirty = False
  554. if not pending:
  555. break
  556. f = pending.pop()
  557. try:
  558. # gc is disabled, so (assuming no daemonic
  559. # threads) the following is the only line in
  560. # this function which might trigger creation
  561. # of a new finalizer
  562. f()
  563. except Exception:
  564. sys.excepthook(*sys.exc_info())
  565. assert f not in cls._registry
  566. finally:
  567. # prevent any more finalizers from executing during shutdown
  568. finalize._shutdown = True
  569. if reenable_gc:
  570. gc.enable()