• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""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()
665