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