Add support for weak references to code objects. This will be used by an…

Add support for weak references to code objects. This will be used by an optimization in the incoming Python 3 JIT.

Patch by Reid Kleckner!

Doc/library/weakref.rst

@@ -59,13 +59,13 @@ Not all objects can be weakly referenced; those objects which can include class
 instances, functions written in Python (but not in C), methods (both bound and
 unbound), sets, frozensets, file objects, :term:`generator`\s, type objects,
 :class:`DBcursor` objects from the :mod:`bsddb` module, sockets, arrays, deques,
-and regular expression pattern objects.
+regular expression pattern objects, and code objects.
 
 .. versionchanged:: 2.4
    Added support for files, sockets, arrays, and patterns.
 
 .. versionchanged:: 2.7
-   Added support for thread.lock and threading.Lock.
+   Added support for thread.lock, threading.Lock, and code objects.
 
 Several built-in types such as :class:`list` and :class:`dict` do not directly
 support weak references but can add support through subclassing::

Include/code.h

@@ -26,6 +26,7 @@ typedef struct {
     PyObject *co_lnotab;	/* string (encoding addr<->lineno mapping) See
 				   Objects/lnotab_notes.txt for details. */
     void *co_zombieframe;     /* for optimization only (see frameobject.c) */
+    PyObject *co_weakreflist;   /* to support weakrefs to code objects */
 } PyCodeObject;
 
 /* Masks for co_flags above */

Lib/test/test_code.py

@@ -81,8 +81,10 @@ consts: ("'doc string'", 'None')
 """
 
 import unittest
+import weakref
 import _testcapi
 
+
 def consts(t):
     """Yield a doctest-safe sequence of object reprs."""
     for elt in t:
@@ -109,12 +111,37 @@ class CodeTest(unittest.TestCase):
         self.assertEquals(co.co_firstlineno, 15)
 
 
+class CodeWeakRefTest(unittest.TestCase):
+
+    def test_basic(self):
+        # Create a code object in a clean environment so that we know we have
+        # the only reference to it left.
+        namespace = {}
+        exec "def f(): pass" in globals(), namespace
+        f = namespace["f"]
+        del namespace
+
+        self.called = False
+        def callback(code):
+            self.called = True
+
+        # f is now the last reference to the function, and through it, the code
+        # object.  While we hold it, check that we can create a weakref and
+        # deref it.  Then delete it, and check that the callback gets called and
+        # the reference dies.
+        coderef = weakref.ref(f.__code__, callback)
+        self.assertTrue(bool(coderef()))
+        del f
+        self.assertFalse(bool(coderef()))
+        self.assertTrue(self.called)
+
+
 def test_main(verbose=None):
     from test.test_support import run_doctest, run_unittest
     from test import test_code
     run_doctest(test_code, verbose)
-    run_unittest(CodeTest)
+    run_unittest(CodeTest, CodeWeakRefTest)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     test_main()

Lib/test/test_sys.py

@@ -552,7 +552,7 @@ class SizeofTest(unittest.TestCase):
         # complex
         check(complex(0,1), size(h + '2d'))
         # code
-        check(get_cell().func_code, size(h + '4i8Pi2P'))
+        check(get_cell().func_code, size(h + '4i8Pi3P'))
         # BaseException
         check(BaseException(), size(h + '3P'))
         # UnicodeEncodeError

Misc/NEWS

@@ -19,6 +19,8 @@ Core and Builtins
   printed and Python exits.  Initialize the GIL before importing the site
   module.
 
+- Code objects now support weak references.
+
 Library
 -------
 

Objects/codeobject.c

@@ -103,6 +103,7 @@ PyCode_New(int argcount, int nlocals, int stacksize, int flags,
 		Py_INCREF(lnotab);
 		co->co_lnotab = lnotab;
                 co->co_zombieframe = NULL;
+		co->co_weakreflist = NULL;
 	}
 	return co;
 }
@@ -314,6 +315,8 @@ code_dealloc(PyCodeObject *co)
 	Py_XDECREF(co->co_lnotab);
         if (co->co_zombieframe != NULL)
                 PyObject_GC_Del(co->co_zombieframe);
+	if (co->co_weakreflist != NULL)
+		PyObject_ClearWeakRefs((PyObject*)co);
 	PyObject_DEL(co);
 }
 
@@ -490,8 +493,8 @@ PyTypeObject PyCode_Type = {
 	code_doc,			/* tp_doc */
 	0,				/* tp_traverse */
 	0,				/* tp_clear */
-	code_richcompare,				/* tp_richcompare */
-	0,				/* tp_weaklistoffset */
+	code_richcompare,		/* tp_richcompare */
+	offsetof(PyCodeObject, co_weakreflist), /* tp_weaklistoffset */
 	0,				/* tp_iter */
 	0,				/* tp_iternext */
 	0,				/* tp_methods */