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.

Doc/lib/libweakref.tex

@@ -68,7 +68,7 @@ Extension types can easily be made to support weak references; see section
 information.
 
 
-\begin{funcdesc}{ref}{object\optional{, callback}}
+\begin{classdesc}{ref}{object\optional{, callback}}
   Return a weak reference to \var{object}.  The original object can be
   retrieved by calling the reference object if the referent is still
   alive; if the referent is no longer alive, calling the reference
@@ -100,7 +100,11 @@ information.
   \var{callback}).  If either referent has been deleted, the
   references are equal only if the reference objects are the same
   object.
-\end{funcdesc}
+
+  \versionchanged[This is now a subclassable type rather than a
+                  factory function; it derives from \class{object}]
+                  {2.4}
+\end{classdesc}
 
 \begin{funcdesc}{proxy}{object\optional{, callback}}
   Return a proxy to \var{object} which uses a weak reference.  This
@@ -236,6 +240,41 @@ become invalidated before the weak reference is called; the
 idiom shown above is safe in threaded applications as well as
 single-threaded applications.
 
+Specialized versions of \class{ref} objects can be created through
+subclassing.  This is used in the implementation of the
+\class{WeakValueDictionary} to reduce the memory overhead for each
+entry in the mapping.  This may be most useful to associate additional
+information with a reference, but could also be used to insert
+additional processing on calls to retrieve the referent.
+
+This example shows how a subclass of \class{ref} can be used to store
+additional information about an object and affect the value that's
+returned when the referent is accessed:
+
+\begin{verbatim}
+import weakref
+
+class ExtendedRef(weakref.ref):
+    def __new__(cls, ob, callback=None, **annotations):
+        weakref.ref.__new__(cls, ob, callback)
+        self.__counter = 0
+
+    def __init__(self, ob, callback=None, **annotations):
+        super(ExtendedRef, self).__init__(ob, callback)
+        for k, v in annotations:
+            setattr(self, k, v)
+
+    def __call__(self):
+        """Return a pair containing the referent and the number of
+        times the reference has been called.
+        """
+        ob = super(ExtendedRef, self)()
+        if ob is not None:
+            self.__counter += 1
+            ob = (ob, self.__counter)
+        return ob
+\end{verbatim}
+
 
 \subsection{Example \label{weakref-example}}
 

Include/weakrefobject.h

@@ -22,11 +22,16 @@ PyAPI_DATA(PyTypeObject) _PyWeakref_RefType;
 PyAPI_DATA(PyTypeObject) _PyWeakref_ProxyType;
 PyAPI_DATA(PyTypeObject) _PyWeakref_CallableProxyType;
 
-#define PyWeakref_CheckRef(op) \
+#define PyWeakref_CheckRef(op) PyObject_TypeCheck(op, &_PyWeakref_RefType)
+#define PyWeakref_CheckRefExact(op) \
         ((op)->ob_type == &_PyWeakref_RefType)
 #define PyWeakref_CheckProxy(op) \
         (((op)->ob_type == &_PyWeakref_ProxyType) || \
          ((op)->ob_type == &_PyWeakref_CallableProxyType))
+
+/* This macro calls PyWeakref_CheckRef() last since that can involve a
+   function call; this makes it more likely that the function call
+   will be avoided. */
 #define PyWeakref_Check(op) \
         (PyWeakref_CheckRef(op) || PyWeakref_CheckProxy(op))
 

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

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

Misc/NEWS

@@ -12,6 +12,11 @@ What's New in Python 2.4 alpha 1?
 Core and builtins
 -----------------
 
+- weakref.ref is now the type object also known as
+  weakref.ReferenceType; it can be subclassed like any other new-style
+  class.  There's less per-entry overhead in WeakValueDictionary
+  objects now (one object instead of three).
+
 - Bug #951851: Python crashed when reading import table of certain
   Windows DLLs.
 

Modules/_weakref.c

@@ -57,26 +57,6 @@ weakref_getweakrefs(PyObject *self, PyObject *object)
 }
 
 
-PyDoc_STRVAR(weakref_ref__doc__,
-"ref(object[, callback]) -- create a weak reference to 'object';\n"
-"when 'object' is finalized, 'callback' will be called and passed\n"
-"a reference to the weak reference object when 'object' is about\n"
-"to be finalized.");
-
-static PyObject *
-weakref_ref(PyObject *self, PyObject *args)
-{
-    PyObject *object;
-    PyObject *callback = NULL;
-    PyObject *result = NULL;
-
-    if (PyArg_UnpackTuple(args, "ref", 1, 2, &object, &callback)) {
-        result = PyWeakref_NewRef(object, callback);
-    }
-    return result;
-}
-
-
 PyDoc_STRVAR(weakref_proxy__doc__,
 "proxy(object[, callback]) -- create a proxy object that weakly\n"
 "references 'object'.  'callback', if given, is called with a\n"
@@ -104,8 +84,6 @@ weakref_functions[] =  {
      weakref_getweakrefs__doc__},
     {"proxy",           weakref_proxy,                  METH_VARARGS,
      weakref_proxy__doc__},
-    {"ref",             weakref_ref,                    METH_VARARGS,
-     weakref_ref__doc__},
     {NULL, NULL, 0, NULL}
 };
 
@@ -119,6 +97,9 @@ init_weakref(void)
                        "Weak-reference support module.");
     if (m != NULL) {
         Py_INCREF(&_PyWeakref_RefType);
+        PyModule_AddObject(m, "ref",
+                           (PyObject *) &_PyWeakref_RefType);
+        Py_INCREF(&_PyWeakref_RefType);
         PyModule_AddObject(m, "ReferenceType",
                            (PyObject *) &_PyWeakref_RefType);
         Py_INCREF(&_PyWeakref_ProxyType);

Objects/object.c

@@ -1802,6 +1802,9 @@ _Py_ReadyTypes(void)
 	if (PyType_Ready(&PyType_Type) < 0)
 		Py_FatalError("Can't initialize 'type'");
 
+	if (PyType_Ready(&_PyWeakref_RefType) < 0)
+		Py_FatalError("Can't initialize 'weakref'");
+
 	if (PyType_Ready(&PyBool_Type) < 0)
 		Py_FatalError("Can't initialize 'bool'");