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.

1 files changed, 180 insertions, 68 deletions

diff --git a/Objects/weakrefobject.c b/Objects/weakrefobject.c
index 575a928..572c224 100644
--- a/Objects/weakrefobject.c
+++ b/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;
-        result->wr_object = ob;
-        Py_XINCREF(callback);
-        result->wr_callback = callback;
+        init_weakref(result, ob, 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);
-    clear_weakref(self);
-    PyObject_GC_Del(self);
+    PyObject_GC_UnTrack(self);
+    clear_weakref((PyWeakReference *) 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);
-            Py_DECREF(callback);
+            if (callback != NULL) {
+                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);
         }