Make weak references subclassable:

- weakref.ref and weakref.ReferenceType will become aliases for each
  other

- weakref.ref will be a modern, new-style class with proper __new__
  and __init__ methods

- weakref.WeakValueDictionary will have a lighter memory footprint,
  using a new weakref.ref subclass to associate the key with the
  value, allowing us to have only a single object of overhead for each
  dictionary entry (currently, there are 3 objects of overhead per
  entry: a weakref to the value, a weakref to the dictionary, and a
  function object used as a weakref callback; the weakref to the
  dictionary could be avoided without this change)

- a new macro, PyWeakref_CheckRefExact(), will be added

- PyWeakref_CheckRef() will check for subclasses of weakref.ref

This closes SF patch #983019.

2 files changed, 100 insertions, 16 deletions

diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index 0a9e97d..31e2c7f 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -623,6 +623,72 @@ class ReferencesTestCase(TestBase):
         finally:
             gc.set_threshold(*thresholds)
 
+
+class SubclassableWeakrefTestCase(unittest.TestCase):
+
+    def test_subclass_refs(self):
+        class MyRef(weakref.ref):
+            def __init__(self, ob, callback=None, value=42):
+                self.value = value
+                super(MyRef, self).__init__(ob, callback)
+            def __call__(self):
+                self.called = True
+                return super(MyRef, self).__call__()
+        o = Object("foo")
+        mr = MyRef(o, value=24)
+        self.assert_(mr() is o)
+        self.assert_(mr.called)
+        self.assertEqual(mr.value, 24)
+        del o
+        self.assert_(mr() is None)
+        self.assert_(mr.called)
+
+    def test_subclass_refs_dont_replace_standard_refs(self):
+        class MyRef(weakref.ref):
+            pass
+        o = Object(42)
+        r1 = MyRef(o)
+        r2 = weakref.ref(o)
+        self.assert_(r1 is not r2)
+        self.assertEqual(weakref.getweakrefs(o), [r2, r1])
+        self.assertEqual(weakref.getweakrefcount(o), 2)
+        r3 = MyRef(o)
+        self.assertEqual(weakref.getweakrefcount(o), 3)
+        refs = weakref.getweakrefs(o)
+        self.assertEqual(len(refs), 3)
+        self.assert_(r2 is refs[0])
+        self.assert_(r1 in refs[1:])
+        self.assert_(r3 in refs[1:])
+
+    def test_subclass_refs_dont_conflate_callbacks(self):
+        class MyRef(weakref.ref):
+            pass
+        o = Object(42)
+        r1 = MyRef(o, id)
+        r2 = MyRef(o, str)
+        self.assert_(r1 is not r2)
+        refs = weakref.getweakrefs(o)
+        self.assert_(r1 in refs)
+        self.assert_(r2 in refs)
+
+    def test_subclass_refs_with_slots(self):
+        class MyRef(weakref.ref):
+            __slots__ = "slot1", "slot2"
+            def __new__(type, ob, callback, slot1, slot2):
+                return weakref.ref.__new__(type, ob, callback)
+            def __init__(self, ob, callback, slot1, slot2):
+                self.slot1 = slot1
+                self.slot2 = slot2
+            def meth(self):
+                return self.slot1 + self.slot2
+        o = Object(42)
+        r = MyRef(o, None, "abc", "def")
+        self.assertEqual(r.slot1, "abc")
+        self.assertEqual(r.slot2, "def")
+        self.assertEqual(r.meth(), "abcdef")
+        self.failIf(hasattr(r, "__dict__"))
+
+
 class Object:
     def __init__(self, arg):
         self.arg = arg
diff --git a/Lib/weakref.py b/Lib/weakref.py
index 510cd7c..cfe9456 100644
--- a/Lib/weakref.py
+++ b/Lib/weakref.py
@@ -42,6 +42,14 @@ class WeakValueDictionary(UserDict.UserDict):
     # objects are unwrapped on the way out, and we always wrap on the
     # way in).
 
+    def __init__(self, *args, **kw):
+        UserDict.UserDict.__init__(self, *args, **kw)
+        def remove(wr, selfref=ref(self)):
+            self = selfref()
+            if self is not None:
+                del self.data[wr.key]
+        self._remove = remove
+
     def __getitem__(self, key):
         o = self.data[key]()
         if o is None:
@@ -53,7 +61,7 @@ class WeakValueDictionary(UserDict.UserDict):
         return "<WeakValueDictionary at %s>" % id(self)
 
     def __setitem__(self, key, value):
-        self.data[key] = ref(value, self.__makeremove(key))
+        self.data[key] = KeyedRef(value, self._remove, key)
 
     def copy(self):
         new = WeakValueDictionary()
@@ -117,7 +125,7 @@ class WeakValueDictionary(UserDict.UserDict):
         try:
             wr = self.data[key]
         except KeyError:
-            self.data[key] = ref(default, self.__makeremove(key))
+            self.data[key] = KeyedRef(default, self._remove, key)
             return default
         else:
             return wr()
@@ -128,7 +136,7 @@ class WeakValueDictionary(UserDict.UserDict):
             if not hasattr(dict, "items"):
                 dict = type({})(dict)
             for key, o in dict.items():
-                d[key] = ref(o, self.__makeremove(key))
+                d[key] = KeyedRef(o, self._remove, key)
         if len(kwargs):
             self.update(kwargs)
 
@@ -140,12 +148,26 @@ class WeakValueDictionary(UserDict.UserDict):
                 L.append(o)
         return L
 
-    def __makeremove(self, key):
-        def remove(o, selfref=ref(self), key=key):
-            self = selfref()
-            if self is not None:
-                del self.data[key]
-        return remove
+
+class KeyedRef(ref):
+    """Specialized reference that includes a key corresponding to the value.
+
+    This is used in the WeakValueDictionary to avoid having to create
+    a function object for each key stored in the mapping.  A shared
+    callback object can use the 'key' attribute of a KeyedRef instead
+    of getting a reference to the key from an enclosing scope.
+
+    """
+
+    __slots__ = "key",
+
+    def __new__(type, ob, callback, key):
+        self = ref.__new__(type, ob, callback)
+        self.key = key
+        return self
+
+    def __init__(self, ob, callback, key):
+        super(KeyedRef,  self).__init__(ob, callback)
 
 
 class WeakKeyDictionary(UserDict.UserDict):
@@ -298,15 +320,11 @@ class WeakValuedValueIterator(BaseIter):
 
 class WeakValuedItemIterator(BaseIter):
     def __init__(self, weakdict):
-        self._next = weakdict.data.iteritems().next
+        self._next = weakdict.data.itervalues().next
 
     def next(self):
         while 1:
-            key, wr = self._next()
+            wr = self._next()
             value = wr()
             if value is not None:
-                return key, value
-
-
-# no longer needed
-del UserDict
+                return wr.key, value