Patch #403985: Add support for weak-keyed dictionaries

3 files changed, 82 insertions, 3 deletions

diff --git a/Lib/test/output/test_weakref b/Lib/test/output/test_weakref
index b3d6f97..ef74e7e 100644
--- a/Lib/test/output/test_weakref
+++ b/Lib/test/output/test_weakref
@@ -15,6 +15,10 @@ Weak Valued Dictionaries
 objects are stored in weak dict
 weak dict test complete
 
+Weak Keyed Dictionaries
+objects are stored in weak dict
+weak key dict test complete
+
 Non-callable Proxy References
 XXX -- tests not written!
 
diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index fe19373..befa70d 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -153,6 +153,24 @@ dict.clear()
 print "weak dict test complete"
 
 print
+print "Weak Keyed Dictionaries"
+
+dict = weakref.mapping(weakkeys=1)
+objects = map(Object, range(10))
+for o in objects:
+    dict[o] = o.arg
+print "objects are stored in weak dict"
+for o in objects:
+    verify(weakref.getweakrefcount(o) == 1,
+           "wrong number of weak references to %r!" % o)
+    verify(o.arg is dict[o],
+           "wrong object returned by weak dict!")
+del objects,o
+verify(len(dict)==0, "deleting the keys did not clear the dictionary")
+print "weak key dict test complete"
+
+
+print
 print "Non-callable Proxy References"
 print "XXX -- tests not written!"
 
diff --git a/Lib/weakref.py b/Lib/weakref.py
index f1222fa..9d5eac0 100644
--- a/Lib/weakref.py
+++ b/Lib/weakref.py
@@ -20,11 +20,14 @@ from _weakref import \
 ProxyTypes = (ProxyType, CallableProxyType)
 
 
-def mapping(dict=None):
-    return WeakDictionary(dict)
+def mapping(dict=None,weakkeys=0):
+    if weakkeys:
+        return WeakKeyDictionary(dict)
+    else:
+        return WeakValueDictionary(dict)
 
 
-class WeakDictionary(UserDict.UserDict):
+class WeakValueDictionary(UserDict.UserDict):
 
     # We inherit the constructor without worrying about the input
     # dictionary; since it uses our .update() method, we get the right
@@ -112,5 +115,59 @@ class WeakDictionary(UserDict.UserDict):
         return L
 
 
+class WeakKeyDictionary(UserDict.UserDict):
+
+    def __init__(self, dict=None):
+        self.data = {}
+        if dict is not None: self.update(dict)
+        def remove(k, data=self.data):
+            del data[k]
+        self._remove = remove
+
+    def __getitem__(self, key):
+        return self.data[ref(key)]
+
+    def __repr__(self):
+        return "<WeakKeyDictionary at %s>" % id(self)
+
+    def __setitem__(self, key, value):
+        self.data[ref(key, self._remove)] = value
+
+    def copy(self):
+        new = WeakKeyDictionary()
+        for key, value in self.data.items():
+            o = key()
+            if o is not None:
+                new[o] = value
+
+    def get(self, key, default):
+        return self.data.get(ref(key),default)
+
+    def items(self):
+        L = []
+        for key, value in self.data.items():
+            o = key()
+            if o is not None:
+                L.append((o, value))
+        return L
+
+    def popitem(self):
+        while 1:
+            key, value = self.data.popitem()
+            o = key()
+            if o is not None:
+                return o, value
+
+    def setdefault(self, key, default):
+        return self.data.setdefault(ref(key, self._remove),default)
+
+    def update(self, dict):
+        d = self.data
+        L = []
+        for key, value in dict.items():
+            L.append(ref(key, self._remove), value)
+        for key, r in L:
+            d[key] = r
+
 # no longer needed
 del UserDict