[pypy-svn] pypy extend-rweakdict: Add "keyclass" to RWeakValueDictionary constructor

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: extend-rweakdict
Changeset: r42577:37634fa88b3c
Date: 2011-03-12 00:23 +0100
http://bitbucket.org/pypy/pypy/changeset/37634fa88b3c/

Log:	Add "keyclass" to RWeakValueDictionary constructor

diff --git a/pypy/rlib/rweakref.py b/pypy/rlib/rweakref.py
--- a/pypy/rlib/rweakref.py
+++ b/pypy/rlib/rweakref.py
@@ -12,14 +12,17 @@
     Only supports string keys.
     """
 
-    def __init__(self, valueclass):
+    def __init__(self, keyclass, valueclass):
         self._dict = weakref.WeakValueDictionary()
+        self._keyclass = keyclass
         self._valueclass = valueclass
 
     def get(self, key):
+        assert isinstance(key, self._keyclass)
         return self._dict.get(key, None)
 
     def set(self, key, value):
+        assert isinstance(key, self._keyclass)
         if value is None:
             self._dict.pop(key, None)
         else:
@@ -94,7 +97,7 @@
 class Entry(extregistry.ExtRegistryEntry):
     _about_ = RWeakValueDictionary
 
-    def compute_result_annotation(self, s_valueclass):
+    def compute_result_annotation(self, s_keyclass, s_valueclass):
         return SomeWeakValueDict(_getclassdef(s_valueclass))
 
     def specialize_call(self, hop):

diff --git a/pypy/rlib/test/test_rweakvaldict.py b/pypy/rlib/test/test_rweakvaldict.py
--- a/pypy/rlib/test/test_rweakvaldict.py
+++ b/pypy/rlib/test/test_rweakvaldict.py
@@ -23,7 +23,7 @@
         assert d.get("hello") is None
         return x1, x3    # x2 dies
     def f():
-        d = RWeakValueDictionary(X)
+        d = RWeakValueDictionary(str, X)
         x1, x3 = g(d)
         rgc.collect(); rgc.collect()
         assert d.get("abc") is x1
@@ -70,7 +70,7 @@
     interpret(make_test(loop=12), [])
 
 def test_rpython_prebuilt():
-    d = RWeakValueDictionary(X)
+    d = RWeakValueDictionary(str, X)
     living = [X() for i in range(8)]
     for i in range(8):
         d.set(str(i), living[i])
@@ -89,13 +89,13 @@
     interpret(f, [])
 
 def test_rpython_merge_RWeakValueDictionary():
-    empty = RWeakValueDictionary(X)
+    empty = RWeakValueDictionary(str, X)
     def f(n):
         x = X()
         if n:
             d = empty
         else:
-            d = RWeakValueDictionary(X)
+            d = RWeakValueDictionary(str, X)
             d.set("a", x)
         return d.get("a") is x
     assert f(0)
@@ -107,7 +107,7 @@
 def test_rpython_merge_RWeakValueDictionary2():
     class A(object):
         def __init__(self):
-            self.d = RWeakValueDictionary(A)
+            self.d = RWeakValueDictionary(str, A)
         def f(self, key):
             a = A()
             self.d.set(key, a)
@@ -126,8 +126,8 @@
 
     def g(x):
         if x:
-            d = RWeakValueDictionary(X)
+            d = RWeakValueDictionary(str, X)
         else:
-            d = RWeakValueDictionary(Y)
+            d = RWeakValueDictionary(str, Y)
         d.set("x", X())
     py.test.raises(Exception, interpret, g, [1])