weakref.py revision 39b17c513ae7b9baecdc8292876683647186fee4 
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.
24
25ProxyTypes = (ProxyType, CallableProxyType)
26
27__all__ = ["ref", "proxy", "getweakrefcount", "getweakrefs",
28           "WeakKeyDictionary", "ReferenceType", "ProxyType",
29           "CallableProxyType", "ProxyTypes", "WeakValueDictionary",
30           "WeakSet", "WeakMethod"]
31
32
33class WeakMethod(ref):
34    """
35    A custom `weakref.ref` subclass which simulates a weak reference to
36    a bound method, working around the lifetime problem of bound methods.
37    """
38
39    __slots__ = "_func_ref", "_meth_type", "_alive", "__weakref__"
40
41    def __new__(cls, meth, callback=None):
42        try:
43            obj = meth.__self__
44            func = meth.__func__
45        except AttributeError:
46            raise TypeError("argument should be a bound method, not {}"
47                            .format(type(meth))) from None
48        def _cb(arg):
49            # The self-weakref trick is needed to avoid creating a reference
50            # cycle.
51            self = self_wr()
52            if self._alive:
53                self._alive = False
54                if callback is not None:
55                    callback(self)
56        self = ref.__new__(cls, obj, _cb)
57        self._func_ref = ref(func, _cb)
58        self._meth_type = type(meth)
59        self._alive = True
60        self_wr = ref(self)
61        return self
62
63    def __call__(self):
64        obj = super().__call__()
65        func = self._func_ref()
66        if obj is None or func is None:
67            return None
68        return self._meth_type(func, obj)
69
70    def __eq__(self, other):
71        if isinstance(other, WeakMethod):
72            if not self._alive or not other._alive:
73                return self is other
74            return ref.__eq__(self, other) and self._func_ref == other._func_ref
75        return False
76
77    def __ne__(self, other):
78        if isinstance(other, WeakMethod):
79            if not self._alive or not other._alive:
80                return self is not other
81            return ref.__ne__(self, other) or self._func_ref != other._func_ref
82        return True
83
84    __hash__ = ref.__hash__
85
86
87class WeakValueDictionary(collections.MutableMapping):
88    """Mapping class that references values weakly.
89
90    Entries in the dictionary will be discarded when no strong
91    reference to the value exists anymore
92    """
93    # We inherit the constructor without worrying about the input
94    # dictionary; since it uses our .update() method, we get the right
95    # checks (if the other dictionary is a WeakValueDictionary,
96    # objects are unwrapped on the way out, and we always wrap on the
97    # way in).
98
99    def __init__(self, *args, **kw):
100        def remove(wr, selfref=ref(self)):
101            self = selfref()
102            if self is not None:
103                if self._iterating:
104                    self._pending_removals.append(wr.key)
105                else:
106                    del self.data[wr.key]
107        self._remove = remove
108        # A list of keys to be removed
109        self._pending_removals = []
110        self._iterating = set()
111        self.data = d = {}
112        self.update(*args, **kw)
113
114    def _commit_removals(self):
115        l = self._pending_removals
116        d = self.data
117        # We shouldn't encounter any KeyError, because this method should
118        # always be called *before* mutating the dict.
119        while l:
120            del d[l.pop()]
121
122    def __getitem__(self, key):
123        o = self.data[key]()
124        if o is None:
125            raise KeyError(key)
126        else:
127            return o
128
129    def __delitem__(self, key):
130        if self._pending_removals:
131            self._commit_removals()
132        del self.data[key]
133
134    def __len__(self):
135        return len(self.data) - len(self._pending_removals)
136
137    def __contains__(self, key):
138        try:
139            o = self.data[key]()
140        except KeyError:
141            return False
142        return o is not None
143
144    def __repr__(self):
145        return "<WeakValueDictionary at %s>" % id(self)
146
147    def __setitem__(self, key, value):
148        if self._pending_removals:
149            self._commit_removals()
150        self.data[key] = KeyedRef(value, self._remove, key)
151
152    def copy(self):
153        new = WeakValueDictionary()
154        for key, wr in self.data.items():
155            o = wr()
156            if o is not None:
157                new[key] = o
158        return new
159
160    __copy__ = copy
161
162    def __deepcopy__(self, memo):
163        from copy import deepcopy
164        new = self.__class__()
165        for key, wr in self.data.items():
166            o = wr()
167            if o is not None:
168                new[deepcopy(key, memo)] = o
169        return new
170
171    def get(self, key, default=None):
172        try:
173            wr = self.data[key]
174        except KeyError:
175            return default
176        else:
177            o = wr()
178            if o is None:
179                # This should only happen
180                return default
181            else:
182                return o
183
184    def items(self):
185        with _IterationGuard(self):
186            for k, wr in self.data.items():
187                v = wr()
188                if v is not None:
189                    yield k, v
190
191    def keys(self):
192        with _IterationGuard(self):
193            for k, wr in self.data.items():
194                if wr() is not None:
195                    yield k
196
197    __iter__ = keys
198
199    def itervaluerefs(self):
200        """Return an iterator that yields the weak references to the values.
201
202        The references are not guaranteed to be 'live' at the time
203        they are used, so the result of calling the references needs
204        to be checked before being used.  This can be used to avoid
205        creating references that will cause the garbage collector to
206        keep the values around longer than needed.
207
208        """
209        with _IterationGuard(self):
210            yield from self.data.values()
211
212    def values(self):
213        with _IterationGuard(self):
214            for wr in self.data.values():
215                obj = wr()
216                if obj is not None:
217                    yield obj
218
219    def popitem(self):
220        if self._pending_removals:
221            self._commit_removals()
222        while True:
223            key, wr = self.data.popitem()
224            o = wr()
225            if o is not None:
226                return key, o
227
228    def pop(self, key, *args):
229        if self._pending_removals:
230            self._commit_removals()
231        try:
232            o = self.data.pop(key)()
233        except KeyError:
234            if args:
235                return args[0]
236            raise
237        if o is None:
238            raise KeyError(key)
239        else:
240            return o
241
242    def setdefault(self, key, default=None):
243        try:
244            wr = self.data[key]
245        except KeyError:
246            if self._pending_removals:
247                self._commit_removals()
248            self.data[key] = KeyedRef(default, self._remove, key)
249            return default
250        else:
251            return wr()
252
253    def update(self, dict=None, **kwargs):
254        if self._pending_removals:
255            self._commit_removals()
256        d = self.data
257        if dict is not None:
258            if not hasattr(dict, "items"):
259                dict = type({})(dict)
260            for key, o in dict.items():
261                d[key] = KeyedRef(o, self._remove, key)
262        if len(kwargs):
263            self.update(kwargs)
264
265    def valuerefs(self):
266        """Return a list of weak references to the values.
267
268        The references are not guaranteed to be 'live' at the time
269        they are used, so the result of calling the references needs
270        to be checked before being used.  This can be used to avoid
271        creating references that will cause the garbage collector to
272        keep the values around longer than needed.
273
274        """
275        return list(self.data.values())
276
277
278class KeyedRef(ref):
279    """Specialized reference that includes a key corresponding to the value.
280
281    This is used in the WeakValueDictionary to avoid having to create
282    a function object for each key stored in the mapping.  A shared
283    callback object can use the 'key' attribute of a KeyedRef instead
284    of getting a reference to the key from an enclosing scope.
285
286    """
287
288    __slots__ = "key",
289
290    def __new__(type, ob, callback, key):
291        self = ref.__new__(type, ob, callback)
292        self.key = key
293        return self
294
295    def __init__(self, ob, callback, key):
296        super().__init__(ob, callback)
297
298
299class WeakKeyDictionary(collections.MutableMapping):
300    """ Mapping class that references keys weakly.
301
302    Entries in the dictionary will be discarded when there is no
303    longer a strong reference to the key. This can be used to
304    associate additional data with an object owned by other parts of
305    an application without adding attributes to those objects. This
306    can be especially useful with objects that override attribute
307    accesses.
308    """
309
310    def __init__(self, dict=None):
311        self.data = {}
312        def remove(k, selfref=ref(self)):
313            self = selfref()
314            if self is not None:
315                if self._iterating:
316                    self._pending_removals.append(k)
317                else:
318                    del self.data[k]
319        self._remove = remove
320        # A list of dead weakrefs (keys to be removed)
321        self._pending_removals = []
322        self._iterating = set()
323        if dict is not None:
324            self.update(dict)
325
326    def _commit_removals(self):
327        # NOTE: We don't need to call this method before mutating the dict,
328        # because a dead weakref never compares equal to a live weakref,
329        # even if they happened to refer to equal objects.
330        # However, it means keys may already have been removed.
331        l = self._pending_removals
332        d = self.data
333        while l:
334            try:
335                del d[l.pop()]
336            except KeyError:
337                pass
338
339    def __delitem__(self, key):
340        del self.data[ref(key)]
341
342    def __getitem__(self, key):
343        return self.data[ref(key)]
344
345    def __len__(self):
346        return len(self.data) - len(self._pending_removals)
347
348    def __repr__(self):
349        return "<WeakKeyDictionary at %s>" % id(self)
350
351    def __setitem__(self, key, value):
352        self.data[ref(key, self._remove)] = value
353
354    def copy(self):
355        new = WeakKeyDictionary()
356        for key, value in self.data.items():
357            o = key()
358            if o is not None:
359                new[o] = value
360        return new
361
362    __copy__ = copy
363
364    def __deepcopy__(self, memo):
365        from copy import deepcopy
366        new = self.__class__()
367        for key, value in self.data.items():
368            o = key()
369            if o is not None:
370                new[o] = deepcopy(value, memo)
371        return new
372
373    def get(self, key, default=None):
374        return self.data.get(ref(key),default)
375
376    def __contains__(self, key):
377        try:
378            wr = ref(key)
379        except TypeError:
380            return False
381        return wr in self.data
382
383    def items(self):
384        with _IterationGuard(self):
385            for wr, value in self.data.items():
386                key = wr()
387                if key is not None:
388                    yield key, value
389
390    def keys(self):
391        with _IterationGuard(self):
392            for wr in self.data:
393                obj = wr()
394                if obj is not None:
395                    yield obj
396
397    __iter__ = keys
398
399    def values(self):
400        with _IterationGuard(self):
401            for wr, value in self.data.items():
402                if wr() is not None:
403                    yield value
404
405    def keyrefs(self):
406        """Return a list of weak references to the keys.
407
408        The references are not guaranteed to be 'live' at the time
409        they are used, so the result of calling the references needs
410        to be checked before being used.  This can be used to avoid
411        creating references that will cause the garbage collector to
412        keep the keys around longer than needed.
413
414        """
415        return list(self.data)
416
417    def popitem(self):
418        while True:
419            key, value = self.data.popitem()
420            o = key()
421            if o is not None:
422                return o, value
423
424    def pop(self, key, *args):
425        return self.data.pop(ref(key), *args)
426
427    def setdefault(self, key, default=None):
428        return self.data.setdefault(ref(key, self._remove),default)
429
430    def update(self, dict=None, **kwargs):
431        d = self.data
432        if dict is not None:
433            if not hasattr(dict, "items"):
434                dict = type({})(dict)
435            for key, value in dict.items():
436                d[ref(key, self._remove)] = value
437        if len(kwargs):
438            self.update(kwargs)
439