Home | History | Annotate | Download | only in Lib
      1 """Weak reference support for Python.
      2 
      3 This module is an implementation of PEP 205:
      4 
      5 http://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 
     12 from _weakref import (
     13      getweakrefcount,
     14      getweakrefs,
     15      ref,
     16      proxy,
     17      CallableProxyType,
     18      ProxyType,
     19      ReferenceType,
     20      _remove_dead_weakref)
     21 
     22 from _weakrefset import WeakSet, _IterationGuard
     23 
     24 import collections  # Import after _weakref to avoid circular import.
     25 import sys
     26 import itertools
     27 
     28 ProxyTypes = (ProxyType, CallableProxyType)
     29 
     30 __all__ = ["ref", "proxy", "getweakrefcount", "getweakrefs",
     31            "WeakKeyDictionary", "ReferenceType", "ProxyType",
     32            "CallableProxyType", "ProxyTypes", "WeakValueDictionary",
     33            "WeakSet", "WeakMethod", "finalize"]
     34 
     35 
     36 class 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 
     90 class 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 
    322 class 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 
    343 class 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 
    498 class 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