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