• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""Weak reference support for Python.
2
3This module is an implementation of PEP 205:
4
5https://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, _atomic_removal=_remove_dead_weakref):
123        pop = self._pending_removals.pop
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 True:
128            try:
129                key = pop()
130            except IndexError:
131                return
132            _atomic_removal(d, key)
133
134    def __getitem__(self, key):
135        if self._pending_removals:
136            self._commit_removals()
137        o = self.data[key]()
138        if o is None:
139            raise KeyError(key)
140        else:
141            return o
142
143    def __delitem__(self, key):
144        if self._pending_removals:
145            self._commit_removals()
146        del self.data[key]
147
148    def __len__(self):
149        if self._pending_removals:
150            self._commit_removals()
151        return len(self.data)
152
153    def __contains__(self, key):
154        if self._pending_removals:
155            self._commit_removals()
156        try:
157            o = self.data[key]()
158        except KeyError:
159            return False
160        return o is not None
161
162    def __repr__(self):
163        return "<%s at %#x>" % (self.__class__.__name__, id(self))
164
165    def __setitem__(self, key, value):
166        if self._pending_removals:
167            self._commit_removals()
168        self.data[key] = KeyedRef(value, self._remove, key)
169
170    def copy(self):
171        if self._pending_removals:
172            self._commit_removals()
173        new = WeakValueDictionary()
174        with _IterationGuard(self):
175            for key, wr in self.data.items():
176                o = wr()
177                if o is not None:
178                    new[key] = o
179        return new
180
181    __copy__ = copy
182
183    def __deepcopy__(self, memo):
184        from copy import deepcopy
185        if self._pending_removals:
186            self._commit_removals()
187        new = self.__class__()
188        with _IterationGuard(self):
189            for key, wr in self.data.items():
190                o = wr()
191                if o is not None:
192                    new[deepcopy(key, memo)] = o
193        return new
194
195    def get(self, key, default=None):
196        if self._pending_removals:
197            self._commit_removals()
198        try:
199            wr = self.data[key]
200        except KeyError:
201            return default
202        else:
203            o = wr()
204            if o is None:
205                # This should only happen
206                return default
207            else:
208                return o
209
210    def items(self):
211        if self._pending_removals:
212            self._commit_removals()
213        with _IterationGuard(self):
214            for k, wr in self.data.items():
215                v = wr()
216                if v is not None:
217                    yield k, v
218
219    def keys(self):
220        if self._pending_removals:
221            self._commit_removals()
222        with _IterationGuard(self):
223            for k, wr in self.data.items():
224                if wr() is not None:
225                    yield k
226
227    __iter__ = keys
228
229    def itervaluerefs(self):
230        """Return an iterator that yields the weak references to the values.
231
232        The references are not guaranteed to be 'live' at the time
233        they are used, so the result of calling the references needs
234        to be checked before being used.  This can be used to avoid
235        creating references that will cause the garbage collector to
236        keep the values around longer than needed.
237
238        """
239        if self._pending_removals:
240            self._commit_removals()
241        with _IterationGuard(self):
242            yield from self.data.values()
243
244    def values(self):
245        if self._pending_removals:
246            self._commit_removals()
247        with _IterationGuard(self):
248            for wr in self.data.values():
249                obj = wr()
250                if obj is not None:
251                    yield obj
252
253    def popitem(self):
254        if self._pending_removals:
255            self._commit_removals()
256        while True:
257            key, wr = self.data.popitem()
258            o = wr()
259            if o is not None:
260                return key, o
261
262    def pop(self, key, *args):
263        if self._pending_removals:
264            self._commit_removals()
265        try:
266            o = self.data.pop(key)()
267        except KeyError:
268            o = None
269        if o is None:
270            if args:
271                return args[0]
272            else:
273                raise KeyError(key)
274        else:
275            return o
276
277    def setdefault(self, key, default=None):
278        try:
279            o = self.data[key]()
280        except KeyError:
281            o = None
282        if o is None:
283            if self._pending_removals:
284                self._commit_removals()
285            self.data[key] = KeyedRef(default, self._remove, key)
286            return default
287        else:
288            return o
289
290    def update(self, other=None, /, **kwargs):
291        if self._pending_removals:
292            self._commit_removals()
293        d = self.data
294        if other is not None:
295            if not hasattr(other, "items"):
296                other = dict(other)
297            for key, o in other.items():
298                d[key] = KeyedRef(o, self._remove, key)
299        for key, o in kwargs.items():
300            d[key] = KeyedRef(o, self._remove, key)
301
302    def valuerefs(self):
303        """Return a list of weak references to the values.
304
305        The references are not guaranteed to be 'live' at the time
306        they are used, so the result of calling the references needs
307        to be checked before being used.  This can be used to avoid
308        creating references that will cause the garbage collector to
309        keep the values around longer than needed.
310
311        """
312        if self._pending_removals:
313            self._commit_removals()
314        return list(self.data.values())
315
316    def __ior__(self, other):
317        self.update(other)
318        return self
319
320    def __or__(self, other):
321        if isinstance(other, _collections_abc.Mapping):
322            c = self.copy()
323            c.update(other)
324            return c
325        return NotImplemented
326
327    def __ror__(self, other):
328        if isinstance(other, _collections_abc.Mapping):
329            c = self.__class__()
330            c.update(other)
331            c.update(self)
332            return c
333        return NotImplemented
334
335
336class KeyedRef(ref):
337    """Specialized reference that includes a key corresponding to the value.
338
339    This is used in the WeakValueDictionary to avoid having to create
340    a function object for each key stored in the mapping.  A shared
341    callback object can use the 'key' attribute of a KeyedRef instead
342    of getting a reference to the key from an enclosing scope.
343
344    """
345
346    __slots__ = "key",
347
348    def __new__(type, ob, callback, key):
349        self = ref.__new__(type, ob, callback)
350        self.key = key
351        return self
352
353    def __init__(self, ob, callback, key):
354        super().__init__(ob, callback)
355
356
357class WeakKeyDictionary(_collections_abc.MutableMapping):
358    """ Mapping class that references keys weakly.
359
360    Entries in the dictionary will be discarded when there is no
361    longer a strong reference to the key. This can be used to
362    associate additional data with an object owned by other parts of
363    an application without adding attributes to those objects. This
364    can be especially useful with objects that override attribute
365    accesses.
366    """
367
368    def __init__(self, dict=None):
369        self.data = {}
370        def remove(k, selfref=ref(self)):
371            self = selfref()
372            if self is not None:
373                if self._iterating:
374                    self._pending_removals.append(k)
375                else:
376                    try:
377                        del self.data[k]
378                    except KeyError:
379                        pass
380        self._remove = remove
381        # A list of dead weakrefs (keys to be removed)
382        self._pending_removals = []
383        self._iterating = set()
384        self._dirty_len = False
385        if dict is not None:
386            self.update(dict)
387
388    def _commit_removals(self):
389        # NOTE: We don't need to call this method before mutating the dict,
390        # because a dead weakref never compares equal to a live weakref,
391        # even if they happened to refer to equal objects.
392        # However, it means keys may already have been removed.
393        pop = self._pending_removals.pop
394        d = self.data
395        while True:
396            try:
397                key = pop()
398            except IndexError:
399                return
400
401            try:
402                del d[key]
403            except KeyError:
404                pass
405
406    def _scrub_removals(self):
407        d = self.data
408        self._pending_removals = [k for k in self._pending_removals if k in d]
409        self._dirty_len = False
410
411    def __delitem__(self, key):
412        self._dirty_len = True
413        del self.data[ref(key)]
414
415    def __getitem__(self, key):
416        return self.data[ref(key)]
417
418    def __len__(self):
419        if self._dirty_len and self._pending_removals:
420            # self._pending_removals may still contain keys which were
421            # explicitly removed, we have to scrub them (see issue #21173).
422            self._scrub_removals()
423        return len(self.data) - len(self._pending_removals)
424
425    def __repr__(self):
426        return "<%s at %#x>" % (self.__class__.__name__, id(self))
427
428    def __setitem__(self, key, value):
429        self.data[ref(key, self._remove)] = value
430
431    def copy(self):
432        new = WeakKeyDictionary()
433        with _IterationGuard(self):
434            for key, value in self.data.items():
435                o = key()
436                if o is not None:
437                    new[o] = value
438        return new
439
440    __copy__ = copy
441
442    def __deepcopy__(self, memo):
443        from copy import deepcopy
444        new = self.__class__()
445        with _IterationGuard(self):
446            for key, value in self.data.items():
447                o = key()
448                if o is not None:
449                    new[o] = deepcopy(value, memo)
450        return new
451
452    def get(self, key, default=None):
453        return self.data.get(ref(key),default)
454
455    def __contains__(self, key):
456        try:
457            wr = ref(key)
458        except TypeError:
459            return False
460        return wr in self.data
461
462    def items(self):
463        with _IterationGuard(self):
464            for wr, value in self.data.items():
465                key = wr()
466                if key is not None:
467                    yield key, value
468
469    def keys(self):
470        with _IterationGuard(self):
471            for wr in self.data:
472                obj = wr()
473                if obj is not None:
474                    yield obj
475
476    __iter__ = keys
477
478    def values(self):
479        with _IterationGuard(self):
480            for wr, value in self.data.items():
481                if wr() is not None:
482                    yield value
483
484    def keyrefs(self):
485        """Return a list of weak references to the keys.
486
487        The references are not guaranteed to be 'live' at the time
488        they are used, so the result of calling the references needs
489        to be checked before being used.  This can be used to avoid
490        creating references that will cause the garbage collector to
491        keep the keys around longer than needed.
492
493        """
494        return list(self.data)
495
496    def popitem(self):
497        self._dirty_len = True
498        while True:
499            key, value = self.data.popitem()
500            o = key()
501            if o is not None:
502                return o, value
503
504    def pop(self, key, *args):
505        self._dirty_len = True
506        return self.data.pop(ref(key), *args)
507
508    def setdefault(self, key, default=None):
509        return self.data.setdefault(ref(key, self._remove),default)
510
511    def update(self, dict=None, /, **kwargs):
512        d = self.data
513        if dict is not None:
514            if not hasattr(dict, "items"):
515                dict = type({})(dict)
516            for key, value in dict.items():
517                d[ref(key, self._remove)] = value
518        if len(kwargs):
519            self.update(kwargs)
520
521    def __ior__(self, other):
522        self.update(other)
523        return self
524
525    def __or__(self, other):
526        if isinstance(other, _collections_abc.Mapping):
527            c = self.copy()
528            c.update(other)
529            return c
530        return NotImplemented
531
532    def __ror__(self, other):
533        if isinstance(other, _collections_abc.Mapping):
534            c = self.__class__()
535            c.update(other)
536            c.update(self)
537            return c
538        return NotImplemented
539
540
541class finalize:
542    """Class for finalization of weakrefable objects
543
544    finalize(obj, func, *args, **kwargs) returns a callable finalizer
545    object which will be called when obj is garbage collected. The
546    first time the finalizer is called it evaluates func(*arg, **kwargs)
547    and returns the result. After this the finalizer is dead, and
548    calling it just returns None.
549
550    When the program exits any remaining finalizers for which the
551    atexit attribute is true will be run in reverse order of creation.
552    By default atexit is true.
553    """
554
555    # Finalizer objects don't have any state of their own.  They are
556    # just used as keys to lookup _Info objects in the registry.  This
557    # ensures that they cannot be part of a ref-cycle.
558
559    __slots__ = ()
560    _registry = {}
561    _shutdown = False
562    _index_iter = itertools.count()
563    _dirty = False
564    _registered_with_atexit = False
565
566    class _Info:
567        __slots__ = ("weakref", "func", "args", "kwargs", "atexit", "index")
568
569    def __init__(self, obj, func, /, *args, **kwargs):
570        if not self._registered_with_atexit:
571            # We may register the exit function more than once because
572            # of a thread race, but that is harmless
573            import atexit
574            atexit.register(self._exitfunc)
575            finalize._registered_with_atexit = True
576        info = self._Info()
577        info.weakref = ref(obj, self)
578        info.func = func
579        info.args = args
580        info.kwargs = kwargs or None
581        info.atexit = True
582        info.index = next(self._index_iter)
583        self._registry[self] = info
584        finalize._dirty = True
585
586    def __call__(self, _=None):
587        """If alive then mark as dead and return func(*args, **kwargs);
588        otherwise return None"""
589        info = self._registry.pop(self, None)
590        if info and not self._shutdown:
591            return info.func(*info.args, **(info.kwargs or {}))
592
593    def detach(self):
594        """If alive then mark as dead and return (obj, func, args, kwargs);
595        otherwise return None"""
596        info = self._registry.get(self)
597        obj = info and info.weakref()
598        if obj is not None and self._registry.pop(self, None):
599            return (obj, info.func, info.args, info.kwargs or {})
600
601    def peek(self):
602        """If alive then return (obj, func, args, kwargs);
603        otherwise return None"""
604        info = self._registry.get(self)
605        obj = info and info.weakref()
606        if obj is not None:
607            return (obj, info.func, info.args, info.kwargs or {})
608
609    @property
610    def alive(self):
611        """Whether finalizer is alive"""
612        return self in self._registry
613
614    @property
615    def atexit(self):
616        """Whether finalizer should be called at exit"""
617        info = self._registry.get(self)
618        return bool(info) and info.atexit
619
620    @atexit.setter
621    def atexit(self, value):
622        info = self._registry.get(self)
623        if info:
624            info.atexit = bool(value)
625
626    def __repr__(self):
627        info = self._registry.get(self)
628        obj = info and info.weakref()
629        if obj is None:
630            return '<%s object at %#x; dead>' % (type(self).__name__, id(self))
631        else:
632            return '<%s object at %#x; for %r at %#x>' % \
633                (type(self).__name__, id(self), type(obj).__name__, id(obj))
634
635    @classmethod
636    def _select_for_exit(cls):
637        # Return live finalizers marked for exit, oldest first
638        L = [(f,i) for (f,i) in cls._registry.items() if i.atexit]
639        L.sort(key=lambda item:item[1].index)
640        return [f for (f,i) in L]
641
642    @classmethod
643    def _exitfunc(cls):
644        # At shutdown invoke finalizers for which atexit is true.
645        # This is called once all other non-daemonic threads have been
646        # joined.
647        reenable_gc = False
648        try:
649            if cls._registry:
650                import gc
651                if gc.isenabled():
652                    reenable_gc = True
653                    gc.disable()
654                pending = None
655                while True:
656                    if pending is None or finalize._dirty:
657                        pending = cls._select_for_exit()
658                        finalize._dirty = False
659                    if not pending:
660                        break
661                    f = pending.pop()
662                    try:
663                        # gc is disabled, so (assuming no daemonic
664                        # threads) the following is the only line in
665                        # this function which might trigger creation
666                        # of a new finalizer
667                        f()
668                    except Exception:
669                        sys.excepthook(*sys.exc_info())
670                    assert f not in cls._registry
671        finally:
672            # prevent any more finalizers from executing during shutdown
673            finalize._shutdown = True
674            if reenable_gc:
675                gc.enable()
676