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