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
23from _weakrefset import WeakSet
24
25from exceptions import ReferenceError
26
27
28ProxyTypes = (ProxyType, CallableProxyType)
29
30__all__ = ["ref", "proxy", "getweakrefcount", "getweakrefs",
31           "WeakKeyDictionary", "ReferenceError", "ReferenceType", "ProxyType",
32           "CallableProxyType", "ProxyTypes", "WeakValueDictionary", 'WeakSet']
33
34
35class WeakValueDictionary(UserDict.UserDict):
36    """Mapping class that references values weakly.
37
38    Entries in the dictionary will be discarded when no strong
39    reference to the value exists anymore
40    """
41    # We inherit the constructor without worrying about the input
42    # dictionary; since it uses our .update() method, we get the right
43    # checks (if the other dictionary is a WeakValueDictionary,
44    # objects are unwrapped on the way out, and we always wrap on the
45    # way in).
46
47    def __init__(self, *args, **kw):
48        def remove(wr, selfref=ref(self)):
49            self = selfref()
50            if self is not None:
51                del self.data[wr.key]
52        self._remove = remove
53        UserDict.UserDict.__init__(self, *args, **kw)
54
55    def __getitem__(self, key):
56        o = self.data[key]()
57        if o is None:
58            raise KeyError, key
59        else:
60            return o
61
62    def __contains__(self, key):
63        try:
64            o = self.data[key]()
65        except KeyError:
66            return False
67        return o is not None
68
69    def has_key(self, key):
70        try:
71            o = self.data[key]()
72        except KeyError:
73            return False
74        return o is not None
75
76    def __repr__(self):
77        return "<WeakValueDictionary at %s>" % id(self)
78
79    def __setitem__(self, key, value):
80        self.data[key] = KeyedRef(value, self._remove, key)
81
82    def copy(self):
83        new = WeakValueDictionary()
84        for key, wr in self.data.items():
85            o = wr()
86            if o is not None:
87                new[key] = o
88        return new
89
90    __copy__ = copy
91
92    def __deepcopy__(self, memo):
93        from copy import deepcopy
94        new = self.__class__()
95        for key, wr in self.data.items():
96            o = wr()
97            if o is not None:
98                new[deepcopy(key, memo)] = o
99        return new
100
101    def get(self, key, default=None):
102        try:
103            wr = self.data[key]
104        except KeyError:
105            return default
106        else:
107            o = wr()
108            if o is None:
109                # This should only happen
110                return default
111            else:
112                return o
113
114    def items(self):
115        L = []
116        for key, wr in self.data.items():
117            o = wr()
118            if o is not None:
119                L.append((key, o))
120        return L
121
122    def iteritems(self):
123        for wr in self.data.itervalues():
124            value = wr()
125            if value is not None:
126                yield wr.key, value
127
128    def iterkeys(self):
129        return self.data.iterkeys()
130
131    def __iter__(self):
132        return self.data.iterkeys()
133
134    def itervaluerefs(self):
135        """Return an iterator that yields the weak references to the values.
136
137        The references are not guaranteed to be 'live' at the time
138        they are used, so the result of calling the references needs
139        to be checked before being used.  This can be used to avoid
140        creating references that will cause the garbage collector to
141        keep the values around longer than needed.
142
143        """
144        return self.data.itervalues()
145
146    def itervalues(self):
147        for wr in self.data.itervalues():
148            obj = wr()
149            if obj is not None:
150                yield obj
151
152    def popitem(self):
153        while 1:
154            key, wr = self.data.popitem()
155            o = wr()
156            if o is not None:
157                return key, o
158
159    def pop(self, key, *args):
160        try:
161            o = self.data.pop(key)()
162        except KeyError:
163            if args:
164                return args[0]
165            raise
166        if o is None:
167            raise KeyError, key
168        else:
169            return o
170
171    def setdefault(self, key, default=None):
172        try:
173            wr = self.data[key]
174        except KeyError:
175            self.data[key] = KeyedRef(default, self._remove, key)
176            return default
177        else:
178            return wr()
179
180    def update(self, dict=None, **kwargs):
181        d = self.data
182        if dict is not None:
183            if not hasattr(dict, "items"):
184                dict = type({})(dict)
185            for key, o in dict.items():
186                d[key] = KeyedRef(o, self._remove, key)
187        if len(kwargs):
188            self.update(kwargs)
189
190    def valuerefs(self):
191        """Return a list of weak references to the values.
192
193        The references are not guaranteed to be 'live' at the time
194        they are used, so the result of calling the references needs
195        to be checked before being used.  This can be used to avoid
196        creating references that will cause the garbage collector to
197        keep the values around longer than needed.
198
199        """
200        return self.data.values()
201
202    def values(self):
203        L = []
204        for wr in self.data.values():
205            o = wr()
206            if o is not None:
207                L.append(o)
208        return L
209
210
211class KeyedRef(ref):
212    """Specialized reference that includes a key corresponding to the value.
213
214    This is used in the WeakValueDictionary to avoid having to create
215    a function object for each key stored in the mapping.  A shared
216    callback object can use the 'key' attribute of a KeyedRef instead
217    of getting a reference to the key from an enclosing scope.
218
219    """
220
221    __slots__ = "key",
222
223    def __new__(type, ob, callback, key):
224        self = ref.__new__(type, ob, callback)
225        self.key = key
226        return self
227
228    def __init__(self, ob, callback, key):
229        super(KeyedRef,  self).__init__(ob, callback)
230
231
232class WeakKeyDictionary(UserDict.UserDict):
233    """ Mapping class that references keys weakly.
234
235    Entries in the dictionary will be discarded when there is no
236    longer a strong reference to the key. This can be used to
237    associate additional data with an object owned by other parts of
238    an application without adding attributes to those objects. This
239    can be especially useful with objects that override attribute
240    accesses.
241    """
242
243    def __init__(self, dict=None):
244        self.data = {}
245        def remove(k, selfref=ref(self)):
246            self = selfref()
247            if self is not None:
248                del self.data[k]
249        self._remove = remove
250        if dict is not None: self.update(dict)
251
252    def __delitem__(self, key):
253        del self.data[ref(key)]
254
255    def __getitem__(self, key):
256        return self.data[ref(key)]
257
258    def __repr__(self):
259        return "<WeakKeyDictionary at %s>" % id(self)
260
261    def __setitem__(self, key, value):
262        self.data[ref(key, self._remove)] = value
263
264    def copy(self):
265        new = WeakKeyDictionary()
266        for key, value in self.data.items():
267            o = key()
268            if o is not None:
269                new[o] = value
270        return new
271
272    __copy__ = copy
273
274    def __deepcopy__(self, memo):
275        from copy import deepcopy
276        new = self.__class__()
277        for key, value in self.data.items():
278            o = key()
279            if o is not None:
280                new[o] = deepcopy(value, memo)
281        return new
282
283    def get(self, key, default=None):
284        return self.data.get(ref(key),default)
285
286    def has_key(self, key):
287        try:
288            wr = ref(key)
289        except TypeError:
290            return 0
291        return wr in self.data
292
293    def __contains__(self, key):
294        try:
295            wr = ref(key)
296        except TypeError:
297            return 0
298        return wr in self.data
299
300    def items(self):
301        L = []
302        for key, value in self.data.items():
303            o = key()
304            if o is not None:
305                L.append((o, value))
306        return L
307
308    def iteritems(self):
309        for wr, value in self.data.iteritems():
310            key = wr()
311            if key is not None:
312                yield key, value
313
314    def iterkeyrefs(self):
315        """Return an iterator that yields the weak references to the keys.
316
317        The references are not guaranteed to be 'live' at the time
318        they are used, so the result of calling the references needs
319        to be checked before being used.  This can be used to avoid
320        creating references that will cause the garbage collector to
321        keep the keys around longer than needed.
322
323        """
324        return self.data.iterkeys()
325
326    def iterkeys(self):
327        for wr in self.data.iterkeys():
328            obj = wr()
329            if obj is not None:
330                yield obj
331
332    def __iter__(self):
333        return self.iterkeys()
334
335    def itervalues(self):
336        return self.data.itervalues()
337
338    def keyrefs(self):
339        """Return a list of weak references to the keys.
340
341        The references are not guaranteed to be 'live' at the time
342        they are used, so the result of calling the references needs
343        to be checked before being used.  This can be used to avoid
344        creating references that will cause the garbage collector to
345        keep the keys around longer than needed.
346
347        """
348        return self.data.keys()
349
350    def keys(self):
351        L = []
352        for wr in self.data.keys():
353            o = wr()
354            if o is not None:
355                L.append(o)
356        return L
357
358    def popitem(self):
359        while 1:
360            key, value = self.data.popitem()
361            o = key()
362            if o is not None:
363                return o, value
364
365    def pop(self, key, *args):
366        return self.data.pop(ref(key), *args)
367
368    def setdefault(self, key, default=None):
369        return self.data.setdefault(ref(key, self._remove),default)
370
371    def update(self, dict=None, **kwargs):
372        d = self.data
373        if dict is not None:
374            if not hasattr(dict, "items"):
375                dict = type({})(dict)
376            for key, value in dict.items():
377                d[ref(key, self._remove)] = value
378        if len(kwargs):
379            self.update(kwargs)
380