SF #1479988: add methods to allow access to weakrefs for the

weakref.WeakKeyDictionary and weakref.WeakValueDictionary

3 files changed, 125 insertions, 0 deletions

diff --git a/Doc/lib/libweakref.tex b/Doc/lib/libweakref.tex
index 840b674..fc949e6 100644
--- a/Doc/lib/libweakref.tex
+++ b/Doc/lib/libweakref.tex
@@ -147,6 +147,24 @@ information.
         to vanish "by magic" (as a side effect of garbage collection).}
 \end{classdesc}
 
+\class{WeakKeyDictionary} objects have the following additional
+methods.  These expose the internal references directly.  The
+references are not guaranteed to be ``live'' at the time they are
+used, so the result of calling the references needs to be checked
+before being used.  This can be used to avoid creating references that
+will cause the garbage collector to keep the keys around longer than
+needed.
+
+\begin{methoddesc}{iterkeyrefs}{}
+  Return an iterator that yields the weak references to the keys.
+  \versionadded{2.5}
+\end{methoddesc}
+
+\begin{methoddesc}{keyrefs}{}
+  Return a list of weak references to the keys.
+  \versionadded{2.5}
+\end{methoddesc}
+
 \begin{classdesc}{WeakValueDictionary}{\optional{dict}}
   Mapping class that references values weakly.  Entries in the
   dictionary will be discarded when no strong reference to the value
@@ -160,6 +178,21 @@ information.
         to vanish "by magic" (as a side effect of garbage collection).}
 \end{classdesc}
 
+\class{WeakValueDictionary} objects have the following additional
+methods.  These method have the same issues as the
+\method{iterkeyrefs()} and \method{keyrefs()} methods of
+\class{WeakKeyDictionary} objects.
+
+\begin{methoddesc}{itervaluerefs}{}
+  Return an iterator that yields the weak references to the values.
+  \versionadded{2.5}
+\end{methoddesc}
+
+\begin{methoddesc}{valuerefs}{}
+  Return a list of weak references to the values.
+  \versionadded{2.5}
+\end{methoddesc}
+
 \begin{datadesc}{ReferenceType}
   The type object for weak references objects.
 \end{datadesc}
diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index 9634ef2..392e5fa 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -769,10 +769,54 @@ class MappingTestCase(TestBase):
         dict, objects = self.make_weak_keyed_dict()
         self.check_iters(dict)
 
+        # Test keyrefs()
+        refs = dict.keyrefs()
+        self.assertEqual(len(refs), len(objects))
+        objects2 = list(objects)
+        for wr in refs:
+            ob = wr()
+            self.assert_(dict.has_key(ob))
+            self.assert_(ob in dict)
+            self.assertEqual(ob.arg, dict[ob])
+            objects2.remove(ob)
+        self.assertEqual(len(objects2), 0)
+
+        # Test iterkeyrefs()
+        objects2 = list(objects)
+        self.assertEqual(len(list(dict.iterkeyrefs())), len(objects))
+        for wr in dict.iterkeyrefs():
+            ob = wr()
+            self.assert_(dict.has_key(ob))
+            self.assert_(ob in dict)
+            self.assertEqual(ob.arg, dict[ob])
+            objects2.remove(ob)
+        self.assertEqual(len(objects2), 0)
+
     def test_weak_valued_iters(self):
         dict, objects = self.make_weak_valued_dict()
         self.check_iters(dict)
 
+        # Test valuerefs()
+        refs = dict.valuerefs()
+        self.assertEqual(len(refs), len(objects))
+        objects2 = list(objects)
+        for wr in refs:
+            ob = wr()
+            self.assertEqual(ob, dict[ob.arg])
+            self.assertEqual(ob.arg, dict[ob.arg].arg)
+            objects2.remove(ob)
+        self.assertEqual(len(objects2), 0)
+
+        # Test itervaluerefs()
+        objects2 = list(objects)
+        self.assertEqual(len(list(dict.itervaluerefs())), len(objects))
+        for wr in dict.itervaluerefs():
+            ob = wr()
+            self.assertEqual(ob, dict[ob.arg])
+            self.assertEqual(ob.arg, dict[ob.arg].arg)
+            objects2.remove(ob)
+        self.assertEqual(len(objects2), 0)
+
     def check_iters(self, dict):
         # item iterator:
         items = dict.items()
diff --git a/Lib/weakref.py b/Lib/weakref.py
index 09bd0be..4f6d757 100644
--- a/Lib/weakref.py
+++ b/Lib/weakref.py
@@ -118,6 +118,18 @@ class WeakValueDictionary(UserDict.UserDict):
     def __iter__(self):
         return self.data.iterkeys()
 
+    def itervaluerefs(self):
+        """Return an iterator that yields the weak references to the values.
+
+        The references are not guaranteed to be 'live' at the time
+        they are used, so the result of calling the references needs
+        to be checked before being used.  This can be used to avoid
+        creating references that will cause the garbage collector to
+        keep the values around longer than needed.
+
+        """
+        return self.data.itervalues()
+
     def itervalues(self):
         for wr in self.data.itervalues():
             obj = wr()
@@ -162,6 +174,18 @@ class WeakValueDictionary(UserDict.UserDict):
         if len(kwargs):
             self.update(kwargs)
 
+    def valuerefs(self):
+        """Return a list of weak references to the values.
+
+        The references are not guaranteed to be 'live' at the time
+        they are used, so the result of calling the references needs
+        to be checked before being used.  This can be used to avoid
+        creating references that will cause the garbage collector to
+        keep the values around longer than needed.
+
+        """
+        return self.data.values()
+
     def values(self):
         L = []
         for wr in self.data.values():
@@ -263,6 +287,18 @@ class WeakKeyDictionary(UserDict.UserDict):
             if key is not None:
                 yield key, value
 
+    def iterkeyrefs(self):
+        """Return an iterator that yields the weak references to the keys.
+
+        The references are not guaranteed to be 'live' at the time
+        they are used, so the result of calling the references needs
+        to be checked before being used.  This can be used to avoid
+        creating references that will cause the garbage collector to
+        keep the keys around longer than needed.
+
+        """
+        return self.data.iterkeys()
+
     def iterkeys(self):
         for wr in self.data.iterkeys():
             obj = wr()
@@ -275,6 +311,18 @@ class WeakKeyDictionary(UserDict.UserDict):
     def itervalues(self):
         return self.data.itervalues()
 
+    def keyrefs(self):
+        """Return a list of weak references to the keys.
+
+        The references are not guaranteed to be 'live' at the time
+        they are used, so the result of calling the references needs
+        to be checked before being used.  This can be used to avoid
+        creating references that will cause the garbage collector to
+        keep the keys around longer than needed.
+
+        """
+        return self.data.keys()
+
     def keys(self):
         L = []
         for wr in self.data.keys():