weakref.py revision a81dd6594067a9501f3ba65b7c05dd50accac5de 
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
21from _weakrefset import WeakSet, _IterationGuard
22
23import collections  # Import after _weakref to avoid circular import.
24import sys
25import itertools
26
27ProxyTypes = (ProxyType, CallableProxyType)
28
29__all__ = ["ref", "proxy", "getweakrefcount", "getweakrefs",
30           "WeakKeyDictionary", "ReferenceType", "ProxyType",
31           "CallableProxyType", "ProxyTypes", "WeakValueDictionary",
32           "WeakSet", "WeakMethod", "finalize"]
33
34
35class WeakMethod(ref):
36    """
37    A custom `weakref.ref` subclass which simulates a weak reference to
38    a bound method, working around the lifetime problem of bound methods.
39    """
40
41    __slots__ = "_func_ref", "_meth_type", "_alive", "__weakref__"
42
43    def __new__(cls, meth, callback=None):
44        try:
45            obj = meth.__self__
46            func = meth.__func__
47        except AttributeError:
48            raise TypeError("argument should be a bound method, not {}"
49                            .format(type(meth))) from None
50        def _cb(arg):
51            # The self-weakref trick is needed to avoid creating a reference
52            # cycle.
53            self = self_wr()
54            if self._alive:
55                self._alive = False
56                if callback is not None:
57                    callback(self)
58        self = ref.__new__(cls, obj, _cb)
59        self._func_ref = ref(func, _cb)
60        self._meth_type = type(meth)
61        self._alive = True
62        self_wr = ref(self)
63        return self
64
65    def __call__(self):
66        obj = super().__call__()
67        func = self._func_ref()
68        if obj is None or func is None:
69            return None
70        return self._meth_type(func, obj)
71
72    def __eq__(self, other):
73        if isinstance(other, WeakMethod):
74            if not self._alive or not other._alive:
75                return self is other
76            return ref.__eq__(self, other) and self._func_ref == other._func_ref
77        return False
78
79    def __ne__(self, other):
80        if isinstance(other, WeakMethod):
81            if not self._alive or not other._alive:
82                return self is not other
83            return ref.__ne__(self, other) or self._func_ref != other._func_ref
84        return True
85
86    __hash__ = ref.__hash__
87
88
89class WeakValueDictionary(collections.MutableMapping):
90    """Mapping class that references values weakly.
91
92    Entries in the dictionary will be discarded when no strong
93    reference to the value exists anymore
94    """
95    # We inherit the constructor without worrying about the input
96    # dictionary; since it uses our .update() method, we get the right
97    # checks (if the other dictionary is a WeakValueDictionary,
98    # objects are unwrapped on the way out, and we always wrap on the
99    # way in).
100
101    def __init__(self, *args, **kw):
102        def remove(wr, selfref=ref(self)):
103            self = selfref()
104            if self is not None:
105                if self._iterating:
106                    self._pending_removals.append(wr.key)
107                else:
108                    del self.data[wr.key]
109        self._remove = remove
110        # A list of keys to be removed
111        self._pending_removals = []
112        self._iterating = set()
113        self.data = d = {}
114        self.update(*args, **kw)
115
116    def _commit_removals(self):
117        l = self._pending_removals
118        d = self.data
119        # We shouldn't encounter any KeyError, because this method should
120        # always be called *before* mutating the dict.
121        while l:
122            del d[l.pop()]
123
124    def __getitem__(self, key):
125        o = self.data[key]()
126        if o is None:
127            raise KeyError(key)
128        else:
129            return o
130
131    def __delitem__(self, key):
132        if self._pending_removals:
133            self._commit_removals()
134        del self.data[key]
135
136    def __len__(self):
137        return len(self.data) - len(self._pending_removals)
138
139    def __contains__(self, key):
140        try:
141            o = self.data[key]()
142        except KeyError:
143            return False
144        return o is not None
145
146    def __repr__(self):
147        return "<WeakValueDictionary at %s>" % id(self)
148
149    def __setitem__(self, key, value):
150        if self._pending_removals:
151            self._commit_removals()
152        self.data[key] = KeyedRef(value, self._remove, key)
153
154    def copy(self):
155        new = WeakValueDictionary()
156        for key, wr in self.data.items():
157            o = wr()
158            if o is not None:
159                new[key] = o
160        return new
161
162    __copy__ = copy
163
164    def __deepcopy__(self, memo):
165        from copy import deepcopy
166        new = self.__class__()
167        for key, wr in self.data.items():
168            o = wr()
169            if o is not None:
170                new[deepcopy(key, memo)] = o
171        return new
172
173    def get(self, key, default=None):
174        try:
175            wr = self.data[key]
176        except KeyError:
177            return default
178        else:
179            o = wr()
180            if o is None:
181                # This should only happen
182                return default
183            else:
184                return o
185
186    def items(self):
187        with _IterationGuard(self):
188            for k, wr in self.data.items():
189                v = wr()
190                if v is not None:
191                    yield k, v
192
193    def keys(self):
194        with _IterationGuard(self):
195            for k, wr in self.data.items():
196                if wr() is not None:
197                    yield k
198
199    __iter__ = keys
200
201    def itervaluerefs(self):
202        """Return an iterator that yields the weak references to the values.
203
204        The references are not guaranteed to be 'live' at the time
205        they are used, so the result of calling the references needs
206        to be checked before being used.  This can be used to avoid
207        creating references that will cause the garbage collector to
208        keep the values around longer than needed.
209
210        """
211        with _IterationGuard(self):
212            yield from self.data.values()
213
214    def values(self):
215        with _IterationGuard(self):
216            for wr in self.data.values():
217                obj = wr()
218                if obj is not None:
219                    yield obj
220
221    def popitem(self):
222        if self._pending_removals:
223            self._commit_removals()
224        while True:
225            key, wr = self.data.popitem()
226            o = wr()
227            if o is not None:
228                return key, o
229
230    def pop(self, key, *args):
231        if self._pending_removals:
232            self._commit_removals()
233        try:
234            o = self.data.pop(key)()
235        except KeyError:
236            if args:
237                return args[0]
238            raise
239        if o is None:
240            raise KeyError(key)
241        else:
242            return o
243
244    def setdefault(self, key, default=None):
245        try:
246            wr = self.data[key]
247        except KeyError:
248            if self._pending_removals:
249                self._commit_removals()
250            self.data[key] = KeyedRef(default, self._remove, key)
251            return default
252        else:
253            return wr()
254
255    def update(self, dict=None, **kwargs):
256        if self._pending_removals:
257            self._commit_removals()
258        d = self.data
259        if dict is not None:
260            if not hasattr(dict, "items"):
261                dict = type({})(dict)
262            for key, o in dict.items():
263                d[key] = KeyedRef(o, self._remove, key)
264        if len(kwargs):
265            self.update(kwargs)
266
267    def valuerefs(self):
268        """Return a list of weak references to the values.
269
270        The references are not guaranteed to be 'live' at the time
271        they are used, so the result of calling the references needs
272        to be checked before being used.  This can be used to avoid
273        creating references that will cause the garbage collector to
274        keep the values around longer than needed.
275
276        """
277        return list(self.data.values())
278
279
280class KeyedRef(ref):
281    """Specialized reference that includes a key corresponding to the value.
282
283    This is used in the WeakValueDictionary to avoid having to create
284    a function object for each key stored in the mapping.  A shared
285    callback object can use the 'key' attribute of a KeyedRef instead
286    of getting a reference to the key from an enclosing scope.
287
288    """
289
290    __slots__ = "key",
291
292    def __new__(type, ob, callback, key):
293        self = ref.__new__(type, ob, callback)
294        self.key = key
295        return self
296
297    def __init__(self, ob, callback, key):
298        super().__init__(ob, callback)
299
300
301class WeakKeyDictionary(collections.MutableMapping):
302    """ Mapping class that references keys weakly.
303
304    Entries in the dictionary will be discarded when there is no
305    longer a strong reference to the key. This can be used to
306    associate additional data with an object owned by other parts of
307    an application without adding attributes to those objects. This
308    can be especially useful with objects that override attribute
309    accesses.
310    """
311
312    def __init__(self, dict=None):
313        self.data = {}
314        def remove(k, selfref=ref(self)):
315            self = selfref()
316            if self is not None:
317                if self._iterating:
318                    self._pending_removals.append(k)
319                else:
320                    del self.data[k]
321        self._remove = remove
322        # A list of dead weakrefs (keys to be removed)
323        self._pending_removals = []
324        self._iterating = set()
325        if dict is not None:
326            self.update(dict)
327
328    def _commit_removals(self):
329        # NOTE: We don't need to call this method before mutating the dict,
330        # because a dead weakref never compares equal to a live weakref,
331        # even if they happened to refer to equal objects.
332        # However, it means keys may already have been removed.
333        l = self._pending_removals
334        d = self.data
335        while l:
336            try:
337                del d[l.pop()]
338            except KeyError:
339                pass
340
341    def __delitem__(self, key):
342        del self.data[ref(key)]
343
344    def __getitem__(self, key):
345        return self.data[ref(key)]
346
347    def __len__(self):
348        return len(self.data) - len(self._pending_removals)
349
350    def __repr__(self):
351        return "<WeakKeyDictionary at %s>" % id(self)
352
353    def __setitem__(self, key, value):
354        self.data[ref(key, self._remove)] = value
355
356    def copy(self):
357        new = WeakKeyDictionary()
358        for key, value in self.data.items():
359            o = key()
360            if o is not None:
361                new[o] = value
362        return new
363
364    __copy__ = copy
365
366    def __deepcopy__(self, memo):
367        from copy import deepcopy
368        new = self.__class__()
369        for key, value in self.data.items():
370            o = key()
371            if o is not None:
372                new[o] = deepcopy(value, memo)
373        return new
374
375    def get(self, key, default=None):
376        return self.data.get(ref(key),default)
377
378    def __contains__(self, key):
379        try:
380            wr = ref(key)
381        except TypeError:
382            return False
383        return wr in self.data
384
385    def items(self):
386        with _IterationGuard(self):
387            for wr, value in self.data.items():
388                key = wr()
389                if key is not None:
390                    yield key, value
391
392    def keys(self):
393        with _IterationGuard(self):
394            for wr in self.data:
395                obj = wr()
396                if obj is not None:
397                    yield obj
398
399    __iter__ = keys
400
401    def values(self):
402        with _IterationGuard(self):
403            for wr, value in self.data.items():
404                if wr() is not None:
405                    yield value
406
407    def keyrefs(self):
408        """Return a list of weak references to the keys.
409
410        The references are not guaranteed to be 'live' at the time
411        they are used, so the result of calling the references needs
412        to be checked before being used.  This can be used to avoid
413        creating references that will cause the garbage collector to
414        keep the keys around longer than needed.
415
416        """
417        return list(self.data)
418
419    def popitem(self):
420        while True:
421            key, value = self.data.popitem()
422            o = key()
423            if o is not None:
424                return o, value
425
426    def pop(self, key, *args):
427        return self.data.pop(ref(key), *args)
428
429    def setdefault(self, key, default=None):
430        return self.data.setdefault(ref(key, self._remove),default)
431
432    def update(self, dict=None, **kwargs):
433        d = self.data
434        if dict is not None:
435            if not hasattr(dict, "items"):
436                dict = type({})(dict)
437            for key, value in dict.items():
438                d[ref(key, self._remove)] = value
439        if len(kwargs):
440            self.update(kwargs)
441
442
443class finalize:
444    """Class for finalization of weakrefable objects
445
446    finalize(obj, func, *args, **kwargs) returns a callable finalizer
447    object which will be called when obj is garbage collected. The
448    first time the finalizer is called it evaluates func(*arg, **kwargs)
449    and returns the result. After this the finalizer is dead, and
450    calling it just returns None.
451
452    When the program exits any remaining finalizers for which the
453    atexit attribute is true will be run in reverse order of creation.
454    By default atexit is true.
455    """
456
457    # Finalizer objects don't have any state of their own.  They are
458    # just used as keys to lookup _Info objects in the registry.  This
459    # ensures that they cannot be part of a ref-cycle.
460
461    __slots__ = ()
462    _registry = {}
463    _shutdown = False
464    _index_iter = itertools.count()
465    _dirty = False
466    _registered_with_atexit = False
467
468    class _Info:
469        __slots__ = ("weakref", "func", "args", "kwargs", "atexit", "index")
470
471    def __init__(self, obj, func, *args, **kwargs):
472        if not self._registered_with_atexit:
473            # We may register the exit function more than once because
474            # of a thread race, but that is harmless
475            import atexit
476            atexit.register(self._exitfunc)
477            finalize._registered_with_atexit = True
478        info = self._Info()
479        info.weakref = ref(obj, self)
480        info.func = func
481        info.args = args
482        info.kwargs = kwargs or None
483        info.atexit = True
484        info.index = next(self._index_iter)
485        self._registry[self] = info
486        finalize._dirty = True
487
488    def __call__(self, _=None):
489        """If alive then mark as dead and return func(*args, **kwargs);
490        otherwise return None"""
491        info = self._registry.pop(self, None)
492        if info and not self._shutdown:
493            return info.func(*info.args, **(info.kwargs or {}))
494
495    def detach(self):
496        """If alive then mark as dead and return (obj, func, args, kwargs);
497        otherwise return None"""
498        info = self._registry.get(self)
499        obj = info and info.weakref()
500        if obj is not None and self._registry.pop(self, None):
501            return (obj, info.func, info.args, info.kwargs or {})
502
503    def peek(self):
504        """If alive then return (obj, func, args, kwargs);
505        otherwise return None"""
506        info = self._registry.get(self)
507        obj = info and info.weakref()
508        if obj is not None:
509            return (obj, info.func, info.args, info.kwargs or {})
510
511    @property
512    def alive(self):
513        """Whether finalizer is alive"""
514        return self in self._registry
515
516    @property
517    def atexit(self):
518        """Whether finalizer should be called at exit"""
519        info = self._registry.get(self)
520        return bool(info) and info.atexit
521
522    @atexit.setter
523    def atexit(self, value):
524        info = self._registry.get(self)
525        if info:
526            info.atexit = bool(value)
527
528    def __repr__(self):
529        info = self._registry.get(self)
530        obj = info and info.weakref()
531        if obj is None:
532            return '<%s object at %#x; dead>' % (type(self).__name__, id(self))
533        else:
534            return '<%s object at %#x; for %r at %#x>' % \
535                (type(self).__name__, id(self), type(obj).__name__, id(obj))
536
537    @classmethod
538    def _select_for_exit(cls):
539        # Return live finalizers marked for exit, oldest first
540        L = [(f,i) for (f,i) in cls._registry.items() if i.atexit]
541        L.sort(key=lambda item:item[1].index)
542        return [f for (f,i) in L]
543
544    @classmethod
545    def _exitfunc(cls):
546        # At shutdown invoke finalizers for which atexit is true.
547        # This is called once all other non-daemonic threads have been
548        # joined.
549        reenable_gc = False
550        try:
551            if cls._registry:
552                import gc
553                if gc.isenabled():
554                    reenable_gc = True
555                    gc.disable()
556                pending = None
557                while True:
558                    if pending is None or finalize._dirty:
559                        pending = cls._select_for_exit()
560                        finalize._dirty = False
561                    if not pending:
562                        break
563                    f = pending.pop()
564                    try:
565                        # gc is disabled, so (assuming no daemonic
566                        # threads) the following is the only line in
567                        # this function which might trigger creation
568                        # of a new finalizer
569                        f()
570                    except Exception:
571                        sys.excepthook(*sys.exc_info())
572                    assert f not in cls._registry
573        finally:
574            # prevent any more finalizers from executing during shutdown
575            finalize._shutdown = True
576            if reenable_gc:
577                gc.enable()
578