[pypy-commit] pypy stmgc-c8: Add _weakref.weakkeyiddict(), a generally useful weak-keyed-dict with

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: stmgc-c8
Changeset: r78252:f5af62298413
Date: 2015-06-23 11:41 +0200
http://bitbucket.org/pypy/pypy/changeset/f5af62298413/

Log:	Add _weakref.weakkeyiddict(), a generally useful weak-keyed-dict
	with identity instead of equality. It is useful to add more
	attributes to existing objects, for example, without attaching them
	to the objects themselves.

diff --git a/pypy/module/_weakref/__init__.py b/pypy/module/_weakref/__init__.py
--- a/pypy/module/_weakref/__init__.py
+++ b/pypy/module/_weakref/__init__.py
@@ -10,5 +10,8 @@
         'ReferenceType': 'interp__weakref.W_Weakref',
         'ProxyType': 'interp__weakref.W_Proxy', 
         'CallableProxyType': 'interp__weakref.W_CallableProxy',
-        'proxy': 'interp__weakref.proxy'
+        'proxy': 'interp__weakref.proxy',
+
+        # PyPy extension
+        'weakkeyiddict': 'weakkeyiddict.W_WeakKeyIdDict',
     }
diff --git a/pypy/module/_weakref/test/test_weakkeyiddict.py b/pypy/module/_weakref/test/test_weakkeyiddict.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/_weakref/test/test_weakkeyiddict.py
@@ -0,0 +1,59 @@
+class AppTestWeakKeyIdDict(object):
+    spaceconfig = dict(usemodules=('_weakref',))
+
+    def test_simple(self):
+        import _weakref
+        class A(object):
+            pass
+        d = _weakref.weakkeyiddict()
+        a1 = A()
+        a2 = A()
+        d[a1] = 11
+        d[a2] = 22.5
+        assert d[a1] == 11
+        assert d[a2] == 22.5
+        assert d.get(a2, 5) == 22.5
+        del d[a2]
+        raises(KeyError, "d[a2]")
+        assert d.get(a2, 5) == 5
+        assert a1 in d
+        assert a2 not in d
+        assert d.setdefault(a1, 82) == 11
+        assert d[a1] == 11
+        assert d.setdefault(a2, 83) == 83
+        assert d[a2] == 83
+
+    def test_nonhashable_key(self):
+        import _weakref
+        d = _weakref.weakkeyiddict()
+        lst = []
+        lst2 = []
+        d[lst] = 84
+        assert lst in d
+        assert lst2 not in d
+        assert d.pop(lst) == 84
+        assert lst not in d
+        assert d.pop(lst, 85) == 85
+
+    def test_collect(self):
+        import _weakref
+        gone = []
+        class A(object):
+            def __del__(self):
+                gone.append(True)
+        d = _weakref.weakkeyiddict()
+        a1 = A()
+        a2 = A()
+        d[a1] = -42
+        d[a2] = 83
+        assert gone == []
+        #
+        del a1
+        tries = 0
+        while not gone:
+            tries += 1
+            if tries > 5:
+                raise AssertionError("a1 doesn't disappear")
+            gc.collect()
+        assert gone == [True]
+        assert d[a2] == 83
diff --git a/pypy/module/_weakref/weakkeyiddict.py b/pypy/module/_weakref/weakkeyiddict.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/_weakref/weakkeyiddict.py
@@ -0,0 +1,82 @@
+#
+# This is _weakref.weakkeyiddict(), a generally useful weak-keyed-dict
+# with identity instead of equality. It is useful to add more
+# attributes to existing objects, for example, without attaching
+# them to the objects themselves.  It can be emulated in pure Python,
+# of course, but given that we already have a class in rlib.rweakref
+# that is doing exactly that in a cheap way, it is far more efficient
+# this way.
+#
+
+from pypy.interpreter.baseobjspace import W_Root
+from pypy.interpreter.typedef import TypeDef
+from pypy.interpreter.gateway import interp2app, unwrap_spec, WrappedDefault
+
+from rpython.rlib.rweakref import RWeakKeyDictionary
+
+
+class W_WeakKeyIdDict(W_Root):
+    def __init__(self):
+        self.data = RWeakKeyDictionary(W_Root, W_Root)
+
+    def getitem_w(self, space, w_key):
+        w_value = self.data.get(w_key)
+        if w_value is None:
+            space.raise_key_error(w_key)
+        return w_value
+
+    def setitem_w(self, space, w_key, w_value):
+        self.data.set(w_key, w_value)
+
+    def delitem_w(self, space, w_key):
+        if self.data.get(w_key) is None:
+            space.raise_key_error(w_key)
+        self.data.set(w_key, None)
+
+    def contains_w(self, space, w_key):
+        return space.wrap(self.data.get(w_key) is not None)
+
+    @unwrap_spec(w_default=WrappedDefault(None))
+    def get_w(self, space, w_key, w_default):
+        w_value = self.data.get(w_key)
+        if w_value is not None:
+            return w_value
+        else:
+            return w_default
+
+    def pop_w(self, space, w_key, w_default=None):
+        w_value = self.data.get(w_key)
+        if w_value is not None:
+            self.data.set(w_key, None)
+            return w_value
+        elif w_default is not None:
+            return w_default
+        else:
+            space.raise_key_error(w_key)
+
+    @unwrap_spec(w_default=WrappedDefault(None))
+    def setdefault_w(self, space, w_key, w_default):
+        w_value = self.data.get(w_key)
+        if w_value is not None:
+            return w_value
+        else:
+            self.data.set(w_key, w_default)
+            return w_default
+
+
+def W_WeakKeyIdDict___new__(space, w_subtype):
+    r = space.allocate_instance(W_WeakKeyIdDict, w_subtype)
+    r.__init__()
+    return space.wrap(r)
+
+W_WeakKeyIdDict.typedef = TypeDef(
+    '_weakref.weakkeyiddict',
+    __new__ = interp2app(W_WeakKeyIdDict___new__),
+    __getitem__ = interp2app(W_WeakKeyIdDict.getitem_w),
+    __setitem__ = interp2app(W_WeakKeyIdDict.setitem_w),
+    __delitem__ = interp2app(W_WeakKeyIdDict.delitem_w),
+    __contains__ = interp2app(W_WeakKeyIdDict.contains_w),
+    get = interp2app(W_WeakKeyIdDict.get_w),
+    pop = interp2app(W_WeakKeyIdDict.pop_w),
+    setdefault = interp2app(W_WeakKeyIdDict.setdefault_w),
+)