[Python-checkins] [3.6] bpo-32176: Set CO_NOFREE in the code object constructor (GH-4684)

[Nick Coghlan]

https://github.com/python/cpython/commit/c8f32aae0aa173e122cf4c0592caec620d0d1de9
commit: c8f32aae0aa173e122cf4c0592caec620d0d1de9
branch: 3.6
author: Nick Coghlan <ncoghlan at gmail.com>
committer: GitHub <noreply at github.com>
date: 2017-12-03T23:32:54+10:00
summary:

[3.6] bpo-32176: Set CO_NOFREE in the code object constructor (GH-4684)

Previously, CO_NOFREE was set in the compiler, which meant
it could end up being set incorrectly when code objects
were created directly. Setting it in the constructor based
on freevars and cellvars ensures it is always accurate,
regardless of how the code object is defined.

(cherry picked from commit 078f1814f1a4413a2a0fdb8cf4490ee0fc98ef34)

files:
A Misc/NEWS.d/next/Core and Builtins/2017-12-02-21-37-22.bpo-32176.Wt25-N.rst
M Lib/test/test_code.py
M Objects/codeobject.c
M Python/compile.c

diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
index 90cb584ac40..55faf4c4279 100644
--- a/Lib/test/test_code.py
+++ b/Lib/test/test_code.py
@@ -102,6 +102,7 @@
 
 """
 
+import inspect
 import sys
 import threading
 import unittest
@@ -130,6 +131,10 @@ def dump(co):
         print("%s: %s" % (attr, getattr(co, "co_" + attr)))
     print("consts:", tuple(consts(co.co_consts)))
 
+# Needed for test_closure_injection below
+# Defined at global scope to avoid implicitly closing over __class__
+def external_getitem(self, i):
+    return f"Foreign getitem: {super().__getitem__(i)}"
 
 class CodeTest(unittest.TestCase):
 
@@ -141,6 +146,46 @@ def test_newempty(self):
         self.assertEqual(co.co_name, "funcname")
         self.assertEqual(co.co_firstlineno, 15)
 
+    @cpython_only
+    def test_closure_injection(self):
+        # From https://bugs.python.org/issue32176
+        from types import FunctionType, CodeType
+
+        def create_closure(__class__):
+            return (lambda: __class__).__closure__
+
+        def new_code(c):
+            '''A new code object with a __class__ cell added to freevars'''
+            return CodeType(
+                c.co_argcount, c.co_kwonlyargcount, c.co_nlocals,
+                c.co_stacksize, c.co_flags, c.co_code, c.co_consts, c.co_names,
+                c.co_varnames, c.co_filename, c.co_name, c.co_firstlineno,
+                c.co_lnotab, c.co_freevars + ('__class__',), c.co_cellvars)
+
+        def add_foreign_method(cls, name, f):
+            code = new_code(f.__code__)
+            assert not f.__closure__
+            closure = create_closure(cls)
+            defaults = f.__defaults__
+            setattr(cls, name, FunctionType(code, globals(), name, defaults, closure))
+
+        class List(list):
+            pass
+
+        add_foreign_method(List, "__getitem__", external_getitem)
+
+        # Ensure the closure injection actually worked
+        function = List.__getitem__
+        class_ref = function.__closure__[0].cell_contents
+        self.assertIs(class_ref, List)
+
+        # Ensure the code correctly indicates it accesses a free variable
+        self.assertFalse(function.__code__.co_flags & inspect.CO_NOFREE,
+                         hex(function.__code__.co_flags))
+
+        # Ensure the zero-arg super() call in the injected method works
+        obj = List([1, 2, 3])
+        self.assertEqual(obj[0], "Foreign getitem: 1")
 
 def isinterned(s):
     return s is sys.intern(('_' + s + '_')[1:-1])
diff --git a/Misc/NEWS.d/next/Core and Builtins/2017-12-02-21-37-22.bpo-32176.Wt25-N.rst b/Misc/NEWS.d/next/Core and Builtins/2017-12-02-21-37-22.bpo-32176.Wt25-N.rst
new file mode 100644
index 00000000000..9d567112866
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2017-12-02-21-37-22.bpo-32176.Wt25-N.rst	
@@ -0,0 +1,5 @@
+co_flags.CO_NOFREE is now always set correctly by the code object
+constructor based on freevars and cellvars, rather than needing to be set
+correctly by the caller. This ensures it will be cleared automatically when
+additional cell references are injected into a modified code object and
+function.
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index d45be4c9679..d1c3abf4861 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -136,12 +136,20 @@ PyCode_New(int argcount, int kwonlyargcount,
     if (PyUnicode_READY(filename) < 0)
         return NULL;
 
-    n_cellvars = PyTuple_GET_SIZE(cellvars);
     intern_strings(names);
     intern_strings(varnames);
     intern_strings(freevars);
     intern_strings(cellvars);
     intern_string_constants(consts);
+
+    /* Check for any inner or outer closure references */
+    n_cellvars = PyTuple_GET_SIZE(cellvars);
+    if (!n_cellvars && !PyTuple_GET_SIZE(freevars)) {
+        flags |= CO_NOFREE;
+    } else {
+        flags &= ~CO_NOFREE;
+    }
+
     /* Create mapping between cells and arguments if needed. */
     if (n_cellvars) {
         Py_ssize_t total_args = argcount + kwonlyargcount +
diff --git a/Python/compile.c b/Python/compile.c
index 797a1840e6a..4d525a02c81 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -5190,18 +5190,6 @@ compute_code_flags(struct compiler *c)
     /* (Only) inherit compilerflags in PyCF_MASK */
     flags |= (c->c_flags->cf_flags & PyCF_MASK);
 
-    n = PyDict_Size(c->u->u_freevars);
-    if (n < 0)
-        return -1;
-    if (n == 0) {
-        n = PyDict_Size(c->u->u_cellvars);
-        if (n < 0)
-            return -1;
-        if (n == 0) {
-            flags |= CO_NOFREE;
-        }
-    }
-
     return flags;
 }