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):
+class KeyedRef(ref):
-            self = selfref()
+    """Specialized reference that includes a key corresponding to the value.
-            if self is not None:
+
-                del self.data[key]
+    This is used in the WeakValueDictionary to avoid having to create
-        return remove
+    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
+                return wr.key, value
-
-
-# no longer needed
-del UserDict

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'");
 

Objects/weakrefobject.c

@@ -19,6 +19,15 @@ _PyWeakref_GetWeakrefCount(PyWeakReference *head)
 }
 
 
+static void
+init_weakref(PyWeakReference *self, PyObject *ob, PyObject *callback)
+{
+    self->hash = -1;
+    self->wr_object = ob;
+    Py_XINCREF(callback);
+    self->wr_callback = callback;
+}
+
 static PyWeakReference *
 new_weakref(PyObject *ob, PyObject *callback)
 {
@@ -26,10 +35,7 @@ new_weakref(PyObject *ob, PyObject *callback)
 
     result = PyObject_GC_New(PyWeakReference, &_PyWeakref_RefType);
     if (result) {
-        result->hash = -1;
+        init_weakref(result, ob, callback);
-        result->wr_object = ob;
-        Py_XINCREF(callback);
-        result->wr_callback = callback;
         PyObject_GC_Track(result);
     }
     return result;
@@ -92,11 +98,11 @@ _PyWeakref_ClearRef(PyWeakReference *self)
 }
 
 static void
-weakref_dealloc(PyWeakReference *self)
+weakref_dealloc(PyObject *self)
 {
-    PyObject_GC_UnTrack((PyObject *)self);
+    PyObject_GC_UnTrack(self);
-    clear_weakref(self);
+    clear_weakref((PyWeakReference *) self);
-    PyObject_GC_Del(self);
+    self->ob_type->tp_free(self);
 }
 
 
@@ -193,6 +199,134 @@ weakref_richcompare(PyWeakReference* self, PyWeakReference* other, int op)
                                 PyWeakref_GET_OBJECT(other), op);
 }
 
+/* Given the head of an object's list of weak references, extract the
+ * two callback-less refs (ref and proxy).  Used to determine if the
+ * shared references exist and to determine the back link for newly
+ * inserted references.
+ */
+static void
+get_basic_refs(PyWeakReference *head,
+               PyWeakReference **refp, PyWeakReference **proxyp)
+{
+    *refp = NULL;
+    *proxyp = NULL;
+
+    if (head != NULL && head->wr_callback == NULL) {
+        /* We need to be careful that the "basic refs" aren't
+           subclasses of the main types.  That complicates this a
+           little. */
+        if (PyWeakref_CheckRefExact(head)) {
+            *refp = head;
+            head = head->wr_next;
+        }
+        if (head != NULL
+            && head->wr_callback == NULL
+            && PyWeakref_CheckProxy(head)) {
+            *proxyp = head;
+            /* head = head->wr_next; */
+        }
+    }
+}
+
+/* Insert 'newref' in the list after 'prev'.  Both must be non-NULL. */
+static void
+insert_after(PyWeakReference *newref, PyWeakReference *prev)
+{
+    newref->wr_prev = prev;
+    newref->wr_next = prev->wr_next;
+    if (prev->wr_next != NULL)
+        prev->wr_next->wr_prev = newref;
+    prev->wr_next = newref;
+}
+
+/* Insert 'newref' at the head of the list; 'list' points to the variable
+ * that stores the head.
+ */
+static void
+insert_head(PyWeakReference *newref, PyWeakReference **list)
+{
+    PyWeakReference *next = *list;
+
+    newref->wr_prev = NULL;
+    newref->wr_next = next;
+    if (next != NULL)
+        next->wr_prev = newref;
+    *list = newref;
+}
+
+static int
+parse_weakref_init_args(char *funcname, PyObject *args, PyObject *kwargs,
+                        PyObject **obp, PyObject **callbackp)
+{
+    /* XXX Should check that kwargs == NULL or is empty. */
+    return PyArg_UnpackTuple(args, funcname, 1, 2, obp, callbackp);
+}
+
+static PyObject *
+weakref___new__(PyTypeObject *type, PyObject *args, PyObject *kwargs)
+{
+    PyWeakReference *self = NULL;
+    PyObject *ob, *callback = NULL;
+
+    if (parse_weakref_init_args("__new__", args, kwargs, &ob, &callback)) {
+        PyWeakReference *ref, *proxy;
+        PyWeakReference **list;
+
+        if (!PyType_SUPPORTS_WEAKREFS(ob->ob_type)) {
+            PyErr_Format(PyExc_TypeError,
+                         "cannot create weak reference to '%s' object",
+                         ob->ob_type->tp_name);
+            return NULL;
+        }
+        if (callback == Py_None)
+            callback = NULL;
+        list = GET_WEAKREFS_LISTPTR(ob);
+        get_basic_refs(*list, &ref, &proxy);
+        if (callback == NULL && type == &_PyWeakref_RefType) {
+            if (ref != NULL) {
+                /* We can re-use an existing reference. */
+                Py_INCREF(ref);
+                return (PyObject *)ref;
+            }
+        }
+        /* We have to create a new reference. */
+        /* Note: the tp_alloc() can trigger cyclic GC, so the weakref
+           list on ob can be mutated.  This means that the ref and
+           proxy pointers we got back earlier may have been collected,
+           so we need to compute these values again before we use
+           them. */
+        self = (PyWeakReference *) (type->tp_alloc(type, 0));
+        if (self != NULL) {
+            init_weakref(self, ob, callback);
+            if (callback == NULL && type == &_PyWeakref_RefType) {
+                insert_head(self, list);
+            }
+            else {
+                PyWeakReference *prev;
+
+                get_basic_refs(*list, &ref, &proxy);
+                prev = (proxy == NULL) ? ref : proxy;
+                if (prev == NULL)
+                    insert_head(self, list);
+                else
+                    insert_after(self, prev);
+            }
+        }
+    }
+    return (PyObject *)self;
+}
+
+static int
+weakref___init__(PyObject *self, PyObject *args, PyObject *kwargs)
+{
+    PyObject *tmp;
+
+    if (parse_weakref_init_args("__init__", args, kwargs, &tmp, &tmp))
+        return 0;
+    else
+        return 1;
+}
+
 
 PyTypeObject
 _PyWeakref_RefType = {
@@ -201,7 +335,7 @@ _PyWeakref_RefType = {
     "weakref",
     sizeof(PyWeakReference),
     0,
-    (destructor)weakref_dealloc,/*tp_dealloc*/
+    weakref_dealloc,            /*tp_dealloc*/
     0,	                        /*tp_print*/
     0,                          /*tp_getattr*/
     0,                          /*tp_setattr*/
@@ -210,18 +344,33 @@ _PyWeakref_RefType = {
     0,                          /*tp_as_number*/
     0,                          /*tp_as_sequence*/
     0,                          /*tp_as_mapping*/
-    (hashfunc)weakref_hash,      /*tp_hash*/
+    (hashfunc)weakref_hash,     /*tp_hash*/
     (ternaryfunc)weakref_call,  /*tp_call*/
     0,                          /*tp_str*/
     0,                          /*tp_getattro*/
     0,                          /*tp_setattro*/
     0,                          /*tp_as_buffer*/
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HAVE_RICHCOMPARE,
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HAVE_RICHCOMPARE
+        | Py_TPFLAGS_BASETYPE,  /*tp_flags*/
     0,                          /*tp_doc*/
     (traverseproc)gc_traverse,  /*tp_traverse*/
     (inquiry)gc_clear,          /*tp_clear*/
     (richcmpfunc)weakref_richcompare,	/*tp_richcompare*/
-    0,				/*tp_weaklistoffset*/
+    0,                          /*tp_weaklistoffset*/
+    0,                          /*tp_iter*/
+    0,                          /*tp_iternext*/
+    0,                          /*tp_methods*/
+    0,                          /*tp_members*/
+    0,                          /*tp_getset*/
+    0,                          /*tp_base*/
+    0,                          /*tp_dict*/
+    0,                          /*tp_descr_get*/
+    0,                          /*tp_descr_set*/
+    0,                          /*tp_dictoffset*/
+    (initproc)weakref___init__, /*tp_init*/
+    PyType_GenericAlloc,        /*tp_alloc*/
+    weakref___new__,            /*tp_new*/
+    PyObject_GC_Del,            /*tp_free*/
 };
 
 
@@ -363,6 +512,15 @@ proxy_nonzero(PyWeakReference *proxy)
         return 1;
 }
 
+static void
+proxy_dealloc(PyWeakReference *self)
+{
+    if (self->wr_callback != NULL)
+        PyObject_GC_UnTrack((PyObject *)self);
+    clear_weakref(self);
+    PyObject_GC_Del(self);
+}
+
 /* sequence slots */
 
 static PyObject *
@@ -496,7 +654,7 @@ _PyWeakref_ProxyType = {
     sizeof(PyWeakReference),
     0,
     /* methods */
-    (destructor)weakref_dealloc,        /* tp_dealloc */
+    (destructor)proxy_dealloc,          /* tp_dealloc */
     0,				        /* tp_print */
     0,				        /* tp_getattr */
     0, 				        /* tp_setattr */
@@ -531,7 +689,7 @@ _PyWeakref_CallableProxyType = {
     sizeof(PyWeakReference),
     0,
     /* methods */
-    (destructor)weakref_dealloc,        /* tp_dealloc */
+    (destructor)proxy_dealloc,          /* tp_dealloc */
     0,				        /* tp_print */
     0,				        /* tp_getattr */
     0, 				        /* tp_setattr */
@@ -558,56 +716,6 @@ _PyWeakref_CallableProxyType = {
 };
 
 
-/* Given the head of an object's list of weak references, extract the
- * two callback-less refs (ref and proxy).  Used to determine if the
- * shared references exist and to determine the back link for newly
- * inserted references.
- */
-static void
-get_basic_refs(PyWeakReference *head,
-               PyWeakReference **refp, PyWeakReference **proxyp)
-{
-    *refp = NULL;
-    *proxyp = NULL;
-
-    if (head != NULL && head->wr_callback == NULL) {
-        if (head->ob_type == &_PyWeakref_RefType) {
-            *refp = head;
-            head = head->wr_next;
-        }
-        if (head != NULL && head->wr_callback == NULL) {
-            *proxyp = head;
-            head = head->wr_next;
-        }
-    }
-}
-
-/* Insert 'newref' in the list after 'prev'.  Both must be non-NULL. */
-static void
-insert_after(PyWeakReference *newref, PyWeakReference *prev)
-{
-    newref->wr_prev = prev;
-    newref->wr_next = prev->wr_next;
-    if (prev->wr_next != NULL)
-        prev->wr_next->wr_prev = newref;
-    prev->wr_next = newref;
-}
-
-/* Insert 'newref' at the head of the list; 'list' points to the variable
- * that stores the head.
- */
-static void
-insert_head(PyWeakReference *newref, PyWeakReference **list)
-{
-    PyWeakReference *next = *list;
-
-    newref->wr_prev = NULL;
-    newref->wr_next = next;
-    if (next != NULL)
-        next->wr_prev = newref;
-    *list = newref;
-}
-
 
 PyObject *
 PyWeakref_NewRef(PyObject *ob, PyObject *callback)
@@ -769,8 +877,10 @@ PyObject_ClearWeakRefs(PyObject *object)
 
             current->wr_callback = NULL;
             clear_weakref(current);
-            handle_callback(current, callback);
+            if (callback != NULL) {
-            Py_DECREF(callback);
+                handle_callback(current, callback);
+                Py_DECREF(callback);
+            }
         }
         else {
             PyObject *tuple = PyTuple_New(count * 2);
@@ -787,10 +897,12 @@ PyObject_ClearWeakRefs(PyObject *object)
                 current = next;
             }
             for (i = 0; i < count; ++i) {
-                PyObject *current = PyTuple_GET_ITEM(tuple, i * 2);
                 PyObject *callback = PyTuple_GET_ITEM(tuple, i * 2 + 1);
 
-                handle_callback((PyWeakReference *)current, callback);
+                if (callback != NULL) {
+                    PyObject *current = PyTuple_GET_ITEM(tuple, i * 2);
+                    handle_callback((PyWeakReference *)current, callback);
+                }
             }
             Py_DECREF(tuple);
         }