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