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