• 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
36class WeakMethod(ref):
37    """
38    A custom `weakref.ref` subclass which simulates a weak reference to
39    a bound method, working around the lifetime problem of bound methods.
40    """
41
42    __slots__ = "_func_ref", "_meth_type", "_alive", "__weakref__"
43
44    def __new__(cls, meth, callback=None):
45        try:
46            obj = meth.__self__
47            func = meth.__func__
48        except AttributeError:
49            raise TypeError("argument should be a bound method, not {}"
50                            .format(type(meth))) from None
51        def _cb(arg):
52            # The self-weakref trick is needed to avoid creating a reference
53            # cycle.
54            self = self_wr()
55            if self._alive:
56                self._alive = False
57                if callback is not None:
58                    callback(self)
59        self = ref.__new__(cls, obj, _cb)
60        self._func_ref = ref(func, _cb)
61        self._meth_type = type(meth)
62        self._alive = True
63        self_wr = ref(self)
64        return self
65
66    def __call__(self):
67        obj = super().__call__()
68        func = self._func_ref()
69        if obj is None or func is None:
70            return None
71        return self._meth_type(func, obj)
72
73    def __eq__(self, other):
74        if isinstance(other, WeakMethod):
75            if not self._alive or not other._alive:
76                return self is other
77            return ref.__eq__(self, other) and self._func_ref == other._func_ref
78        return False
79
80    def __ne__(self, other):
81        if isinstance(other, WeakMethod):
82            if not self._alive or not other._alive:
83                return self is not other
84            return ref.__ne__(self, other) or self._func_ref != other._func_ref
85        return True
86
87    __hash__ = ref.__hash__
88
89
90class WeakValueDictionary(_collections_abc.MutableMapping):
91    """Mapping class that references values weakly.
92
93    Entries in the dictionary will be discarded when no strong
94    reference to the value exists anymore
95    """
96    # We inherit the constructor without worrying about the input
97    # dictionary; since it uses our .update() method, we get the right
98    # checks (if the other dictionary is a WeakValueDictionary,
99    # objects are unwrapped on the way out, and we always wrap on the
100    # way in).
101
102    def __init__(*args, **kw):
103        if not args:
104            raise TypeError("descriptor '__init__' of 'WeakValueDictionary' "
105                            "object needs an argument")
106        self, *args = args
107        if len(args) > 1:
108            raise TypeError('expected at most 1 arguments, got %d' % len(args))
109        def remove(wr, selfref=ref(self), _atomic_removal=_remove_dead_weakref):
110            self = selfref()
111            if self is not None:
112                if self._iterating:
113                    self._pending_removals.append(wr.key)
114                else:
115                    # Atomic removal is necessary since this function
116                    # can be called asynchronously by the GC
117                    _atomic_removal(d, wr.key)
118        self._remove = remove
119        # A list of keys to be removed
120        self._pending_removals = []
121        self._iterating = set()
122        self.data = d = {}
123        self.update(*args, **kw)
124
125    def _commit_removals(self):
126        l = self._pending_removals
127        d = self.data
128        # We shouldn't encounter any KeyError, because this method should
129        # always be called *before* mutating the dict.
130        while l:
131            key = l.pop()
132            _remove_dead_weakref(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(*args, **kwargs):
291        if not args:
292            raise TypeError("descriptor 'update' of 'WeakValueDictionary' "
293                            "object needs an argument")
294        self, *args = args
295        if len(args) > 1:
296            raise TypeError('expected at most 1 arguments, got %d' % len(args))
297        dict = args[0] if args else None
298        if self._pending_removals:
299            self._commit_removals()
300        d = self.data
301        if dict is not None:
302            if not hasattr(dict, "items"):
303                dict = type({})(dict)
304            for key, o in dict.items():
305                d[key] = KeyedRef(o, self._remove, key)
306        if len(kwargs):
307            self.update(kwargs)
308
309    def valuerefs(self):
310        """Return a list of weak references to the values.
311
312        The references are not guaranteed to be 'live' at the time
313        they are used, so the result of calling the references needs
314        to be checked before being used.  This can be used to avoid
315        creating references that will cause the garbage collector to
316        keep the values around longer than needed.
317
318        """
319        if self._pending_removals:
320            self._commit_removals()
321        return list(self.data.values())
322
323
324class KeyedRef(ref):
325    """Specialized reference that includes a key corresponding to the value.
326
327    This is used in the WeakValueDictionary to avoid having to create
328    a function object for each key stored in the mapping.  A shared
329    callback object can use the 'key' attribute of a KeyedRef instead
330    of getting a reference to the key from an enclosing scope.
331
332    """
333
334    __slots__ = "key",
335
336    def __new__(type, ob, callback, key):
337        self = ref.__new__(type, ob, callback)
338        self.key = key
339        return self
340
341    def __init__(self, ob, callback, key):
342        super().__init__(ob, callback)
343
344
345class WeakKeyDictionary(_collections_abc.MutableMapping):
346    """ Mapping class that references keys weakly.
347
348    Entries in the dictionary will be discarded when there is no
349    longer a strong reference to the key. This can be used to
350    associate additional data with an object owned by other parts of
351    an application without adding attributes to those objects. This
352    can be especially useful with objects that override attribute
353    accesses.
354    """
355
356    def __init__(self, dict=None):
357        self.data = {}
358        def remove(k, selfref=ref(self)):
359            self = selfref()
360            if self is not None:
361                if self._iterating:
362                    self._pending_removals.append(k)
363                else:
364                    del self.data[k]
365        self._remove = remove
366        # A list of dead weakrefs (keys to be removed)
367        self._pending_removals = []
368        self._iterating = set()
369        self._dirty_len = False
370        if dict is not None:
371            self.update(dict)
372
373    def _commit_removals(self):
374        # NOTE: We don't need to call this method before mutating the dict,
375        # because a dead weakref never compares equal to a live weakref,
376        # even if they happened to refer to equal objects.
377        # However, it means keys may already have been removed.
378        l = self._pending_removals
379        d = self.data
380        while l:
381            try:
382                del d[l.pop()]
383            except KeyError:
384                pass
385
386    def _scrub_removals(self):
387        d = self.data
388        self._pending_removals = [k for k in self._pending_removals if k in d]
389        self._dirty_len = False
390
391    def __delitem__(self, key):
392        self._dirty_len = True
393        del self.data[ref(key)]
394
395    def __getitem__(self, key):
396        return self.data[ref(key)]
397
398    def __len__(self):
399        if self._dirty_len and self._pending_removals:
400            # self._pending_removals may still contain keys which were
401            # explicitly removed, we have to scrub them (see issue #21173).
402            self._scrub_removals()
403        return len(self.data) - len(self._pending_removals)
404
405    def __repr__(self):
406        return "<%s at %#x>" % (self.__class__.__name__, id(self))
407
408    def __setitem__(self, key, value):
409        self.data[ref(key, self._remove)] = value
410
411    def copy(self):
412        new = WeakKeyDictionary()
413        with _IterationGuard(self):
414            for key, value in self.data.items():
415                o = key()
416                if o is not None:
417                    new[o] = value
418        return new
419
420    __copy__ = copy
421
422    def __deepcopy__(self, memo):
423        from copy import deepcopy
424        new = self.__class__()
425        with _IterationGuard(self):
426            for key, value in self.data.items():
427                o = key()
428                if o is not None:
429                    new[o] = deepcopy(value, memo)
430        return new
431
432    def get(self, key, default=None):
433        return self.data.get(ref(key),default)
434
435    def __contains__(self, key):
436        try:
437            wr = ref(key)
438        except TypeError:
439            return False
440        return wr in self.data
441
442    def items(self):
443        with _IterationGuard(self):
444            for wr, value in self.data.items():
445                key = wr()
446                if key is not None:
447                    yield key, value
448
449    def keys(self):
450        with _IterationGuard(self):
451            for wr in self.data:
452                obj = wr()
453                if obj is not None:
454                    yield obj
455
456    __iter__ = keys
457
458    def values(self):
459        with _IterationGuard(self):
460            for wr, value in self.data.items():
461                if wr() is not None:
462                    yield value
463
464    def keyrefs(self):
465        """Return a list of weak references to the keys.
466
467        The references are not guaranteed to be 'live' at the time
468        they are used, so the result of calling the references needs
469        to be checked before being used.  This can be used to avoid
470        creating references that will cause the garbage collector to
471        keep the keys around longer than needed.
472
473        """
474        return list(self.data)
475
476    def popitem(self):
477        self._dirty_len = True
478        while True:
479            key, value = self.data.popitem()
480            o = key()
481            if o is not None:
482                return o, value
483
484    def pop(self, key, *args):
485        self._dirty_len = True
486        return self.data.pop(ref(key), *args)
487
488    def setdefault(self, key, default=None):
489        return self.data.setdefault(ref(key, self._remove),default)
490
491    def update(self, dict=None, **kwargs):
492        d = self.data
493        if dict is not None:
494            if not hasattr(dict, "items"):
495                dict = type({})(dict)
496            for key, value in dict.items():
497                d[ref(key, self._remove)] = value
498        if len(kwargs):
499            self.update(kwargs)
500
501
502class finalize:
503    """Class for finalization of weakrefable objects
504
505    finalize(obj, func, *args, **kwargs) returns a callable finalizer
506    object which will be called when obj is garbage collected. The
507    first time the finalizer is called it evaluates func(*arg, **kwargs)
508    and returns the result. After this the finalizer is dead, and
509    calling it just returns None.
510
511    When the program exits any remaining finalizers for which the
512    atexit attribute is true will be run in reverse order of creation.
513    By default atexit is true.
514    """
515
516    # Finalizer objects don't have any state of their own.  They are
517    # just used as keys to lookup _Info objects in the registry.  This
518    # ensures that they cannot be part of a ref-cycle.
519
520    __slots__ = ()
521    _registry = {}
522    _shutdown = False
523    _index_iter = itertools.count()
524    _dirty = False
525    _registered_with_atexit = False
526
527    class _Info:
528        __slots__ = ("weakref", "func", "args", "kwargs", "atexit", "index")
529
530    def __init__(self, obj, func, *args, **kwargs):
531        if not self._registered_with_atexit:
532            # We may register the exit function more than once because
533            # of a thread race, but that is harmless
534            import atexit
535            atexit.register(self._exitfunc)
536            finalize._registered_with_atexit = True
537        info = self._Info()
538        info.weakref = ref(obj, self)
539        info.func = func
540        info.args = args
541        info.kwargs = kwargs or None
542        info.atexit = True
543        info.index = next(self._index_iter)
544        self._registry[self] = info
545        finalize._dirty = True
546
547    def __call__(self, _=None):
548        """If alive then mark as dead and return func(*args, **kwargs);
549        otherwise return None"""
550        info = self._registry.pop(self, None)
551        if info and not self._shutdown:
552            return info.func(*info.args, **(info.kwargs or {}))
553
554    def detach(self):
555        """If alive then mark as dead and return (obj, func, args, kwargs);
556        otherwise return None"""
557        info = self._registry.get(self)
558        obj = info and info.weakref()
559        if obj is not None and self._registry.pop(self, None):
560            return (obj, info.func, info.args, info.kwargs or {})
561
562    def peek(self):
563        """If alive then return (obj, func, args, kwargs);
564        otherwise return None"""
565        info = self._registry.get(self)
566        obj = info and info.weakref()
567        if obj is not None:
568            return (obj, info.func, info.args, info.kwargs or {})
569
570    @property
571    def alive(self):
572        """Whether finalizer is alive"""
573        return self in self._registry
574
575    @property
576    def atexit(self):
577        """Whether finalizer should be called at exit"""
578        info = self._registry.get(self)
579        return bool(info) and info.atexit
580
581    @atexit.setter
582    def atexit(self, value):
583        info = self._registry.get(self)
584        if info:
585            info.atexit = bool(value)
586
587    def __repr__(self):
588        info = self._registry.get(self)
589        obj = info and info.weakref()
590        if obj is None:
591            return '<%s object at %#x; dead>' % (type(self).__name__, id(self))
592        else:
593            return '<%s object at %#x; for %r at %#x>' % \
594                (type(self).__name__, id(self), type(obj).__name__, id(obj))
595
596    @classmethod
597    def _select_for_exit(cls):
598        # Return live finalizers marked for exit, oldest first
599        L = [(f,i) for (f,i) in cls._registry.items() if i.atexit]
600        L.sort(key=lambda item:item[1].index)
601        return [f for (f,i) in L]
602
603    @classmethod
604    def _exitfunc(cls):
605        # At shutdown invoke finalizers for which atexit is true.
606        # This is called once all other non-daemonic threads have been
607        # joined.
608        reenable_gc = False
609        try:
610            if cls._registry:
611                import gc
612                if gc.isenabled():
613                    reenable_gc = True
614                    gc.disable()
615                pending = None
616                while True:
617                    if pending is None or finalize._dirty:
618                        pending = cls._select_for_exit()
619                        finalize._dirty = False
620                    if not pending:
621                        break
622                    f = pending.pop()
623                    try:
624                        # gc is disabled, so (assuming no daemonic
625                        # threads) the following is the only line in
626                        # this function which might trigger creation
627                        # of a new finalizer
628                        f()
629                    except Exception:
630                        sys.excepthook(*sys.exc_info())
631                    assert f not in cls._registry
632        finally:
633            # prevent any more finalizers from executing during shutdown
634            finalize._shutdown = True
635            if reenable_gc:
636                gc.enable()
637