bpo-37032: Add CodeType.replace() method (GH-13542)

8 files changed, 401 insertions, 33 deletions

diff --git a/Doc/whatsnew/3.8.rst b/Doc/whatsnew/3.8.rst
index b5d70b5..a94aba6 100644
--- a/Doc/whatsnew/3.8.rst
+++ b/Doc/whatsnew/3.8.rst
@@ -240,6 +240,9 @@ Other Language Changes
   and Windows use this to properly terminate scripts in interactive sessions.
   (Contributed by Google via Gregory P. Smith in :issue:`1054041`.)
 
+* Added new ``replace()`` method to the code type (:class:`types.CodeType`).
+  (Contributed by Victor Stinner in :issue:`37032`.)
+
 
 New Modules
 ===========
@@ -1051,7 +1054,8 @@ Changes in the Python API
 
 * :class:`types.CodeType` has a new parameter in the second position of the
   constructor (*posonlyargcount*) to support positional-only arguments defined
-  in :pep:`570`.
+  in :pep:`570`. A new ``replace()`` method of :class:`types.CodeType` can be
+  used to make the code future-proof.
 
 
 Changes in the C API
diff --git a/Lib/modulefinder.py b/Lib/modulefinder.py
index a36f1b6..e0d2998 100644
--- a/Lib/modulefinder.py
+++ b/Lib/modulefinder.py
@@ -619,13 +619,7 @@ class ModuleFinder:
             if isinstance(consts[i], type(co)):
                 consts[i] = self.replace_paths_in_code(consts[i])
 
-        return types.CodeType(co.co_argcount, co.co_posonlyargcount,
-                              co.co_kwonlyargcount, co.co_nlocals,
-                              co.co_stacksize, co.co_flags,
-                              co.co_code, tuple(consts), co.co_names,
-                              co.co_varnames, new_filename, co.co_name,
-                              co.co_firstlineno, co.co_lnotab, co.co_freevars,
-                              co.co_cellvars)
+        return co.replace(co_consts=tuple(consts), co_filename=new_filename)
 
 
 def test():
diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
index 9bf290d..91008c0 100644
--- a/Lib/test/test_code.py
+++ b/Lib/test/test_code.py
@@ -174,18 +174,14 @@ class CodeTest(unittest.TestCase):
     @cpython_only
     def test_closure_injection(self):
         # From https://bugs.python.org/issue32176
-        from types import FunctionType, CodeType
+        from types import FunctionType
 
         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_posonlyargcount, 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)
+            return c.replace(co_freevars=c.co_freevars + ('__class__',))
 
         def add_foreign_method(cls, name, f):
             code = new_code(f.__code__)
@@ -212,6 +208,64 @@ class CodeTest(unittest.TestCase):
         obj = List([1, 2, 3])
         self.assertEqual(obj[0], "Foreign getitem: 1")
 
+    def test_constructor(self):
+        def func(): pass
+        co = func.__code__
+        CodeType = type(co)
+
+        # test code constructor
+        return CodeType(co.co_argcount,
+                        co.co_posonlyargcount,
+                        co.co_kwonlyargcount,
+                        co.co_nlocals,
+                        co.co_stacksize,
+                        co.co_flags,
+                        co.co_code,
+                        co.co_consts,
+                        co.co_names,
+                        co.co_varnames,
+                        co.co_filename,
+                        co.co_name,
+                        co.co_firstlineno,
+                        co.co_lnotab,
+                        co.co_freevars,
+                        co.co_cellvars)
+
+    def test_replace(self):
+        def func():
+            x = 1
+            return x
+        code = func.__code__
+
+        # different co_name, co_varnames, co_consts
+        def func2():
+            y = 2
+            return y
+        code2 = func.__code__
+
+        for attr, value in (
+            ("co_argcount", 0),
+            ("co_posonlyargcount", 0),
+            ("co_kwonlyargcount", 0),
+            ("co_nlocals", 0),
+            ("co_stacksize", 0),
+            ("co_flags", code.co_flags | inspect.CO_COROUTINE),
+            ("co_firstlineno", 100),
+            ("co_code", code2.co_code),
+            ("co_consts", code2.co_consts),
+            ("co_names", ("myname",)),
+            ("co_varnames", code2.co_varnames),
+            ("co_freevars", ("freevar",)),
+            ("co_cellvars", ("cellvar",)),
+            ("co_filename", "newfilename"),
+            ("co_name", "newname"),
+            ("co_lnotab", code2.co_lnotab),
+        ):
+            with self.subTest(attr=attr, value=value):
+                new_code = code.replace(**{attr: value})
+                self.assertEqual(getattr(new_code, attr), value)
+
+
 def isinterned(s):
     return s is sys.intern(('_' + s + '_')[1:-1])
 
diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py
index 109b014..84cd0da 100644
--- a/Lib/test/test_import/__init__.py
+++ b/Lib/test/test_import/__init__.py
@@ -674,13 +674,7 @@ func_filename = func.__code__.co_filename
         foreign_code = importlib.import_module.__code__
         pos = constants.index(1)
         constants[pos] = foreign_code
-        code = type(code)(code.co_argcount, code.co_posonlyargcount,
-                          code.co_kwonlyargcount,
-                          code.co_nlocals, code.co_stacksize,
-                          code.co_flags, code.co_code, tuple(constants),
-                          code.co_names, code.co_varnames, code.co_filename,
-                          code.co_name, code.co_firstlineno, code.co_lnotab,
-                          code.co_freevars, code.co_cellvars)
+        code = code.replace(co_consts=tuple(constants))
         with open(self.compiled_name, "wb") as f:
             f.write(header)
             marshal.dump(code, f)
diff --git a/Lib/types.py b/Lib/types.py
index 37ba4bb..ea3c0b29 100644
--- a/Lib/types.py
+++ b/Lib/types.py
@@ -262,14 +262,8 @@ def coroutine(func):
         if co_flags & 0x20:
             # TODO: Implement this in C.
             co = func.__code__
-            func.__code__ = CodeType(
-                co.co_argcount, co.co_posonlyargcount, co.co_kwonlyargcount, co.co_nlocals,
-                co.co_stacksize,
-                co.co_flags | 0x100,  # 0x100 == CO_ITERABLE_COROUTINE
-                co.co_code,
-                co.co_consts, co.co_names, co.co_varnames, co.co_filename,
-                co.co_name, co.co_firstlineno, co.co_lnotab, co.co_freevars,
-                co.co_cellvars)
+            # 0x100 == CO_ITERABLE_COROUTINE
+            func.__code__ = co.replace(co_flags=co.co_flags | 0x100)
             return func
 
     # The following code is primarily to support functions that
diff --git a/Misc/NEWS.d/next/Core and Builtins/2019-05-24-12-38-40.bpo-37032.T8rSH8.rst b/Misc/NEWS.d/next/Core and Builtins/2019-05-24-12-38-40.bpo-37032.T8rSH8.rst
new file mode 100644
index 0000000..7e31a84
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2019-05-24-12-38-40.bpo-37032.T8rSH8.rst
@@ -0,0 +1 @@
+Added new ``replace()`` method to the code type (:class:`types.CodeType`).
diff --git a/Objects/clinic/codeobject.c.h b/Objects/clinic/codeobject.c.h
new file mode 100644
index 0000000..ec127ce
--- /dev/null
+++ b/Objects/clinic/codeobject.c.h
@@ -0,0 +1,256 @@
+/*[clinic input]
+preserve
+[clinic start generated code]*/
+
+PyDoc_STRVAR(code_replace__doc__,
+"replace($self, /, *, co_argcount=-1, co_posonlyargcount=-1,\n"
+"        co_kwonlyargcount=-1, co_nlocals=-1, co_stacksize=-1,\n"
+"        co_flags=-1, co_firstlineno=-1, co_code=None, co_consts=None,\n"
+"        co_names=None, co_varnames=None, co_freevars=None,\n"
+"        co_cellvars=None, co_filename=None, co_name=None,\n"
+"        co_lnotab=None)\n"
+"--\n"
+"\n"
+"Return a new code object with new specified fields.");
+
+#define CODE_REPLACE_METHODDEF    \
+    {"replace", (PyCFunction)(void(*)(void))code_replace, METH_FASTCALL|METH_KEYWORDS, code_replace__doc__},
+
+static PyObject *
+code_replace_impl(PyCodeObject *self, int co_argcount,
+                  int co_posonlyargcount, int co_kwonlyargcount,
+                  int co_nlocals, int co_stacksize, int co_flags,
+                  int co_firstlineno, PyBytesObject *co_code,
+                  PyObject *co_consts, PyObject *co_names,
+                  PyObject *co_varnames, PyObject *co_freevars,
+                  PyObject *co_cellvars, PyObject *co_filename,
+                  PyObject *co_name, PyBytesObject *co_lnotab);
+
+static PyObject *
+code_replace(PyCodeObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"co_argcount", "co_posonlyargcount", "co_kwonlyargcount", "co_nlocals", "co_stacksize", "co_flags", "co_firstlineno", "co_code", "co_consts", "co_names", "co_varnames", "co_freevars", "co_cellvars", "co_filename", "co_name", "co_lnotab", NULL};
+    static _PyArg_Parser _parser = {NULL, _keywords, "replace", 0};
+    PyObject *argsbuf[16];
+    Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0;
+    int co_argcount = self->co_argcount;
+    int co_posonlyargcount = self->co_posonlyargcount;
+    int co_kwonlyargcount = self->co_kwonlyargcount;
+    int co_nlocals = self->co_nlocals;
+    int co_stacksize = self->co_stacksize;
+    int co_flags = self->co_flags;
+    int co_firstlineno = self->co_firstlineno;
+    PyBytesObject *co_code = (PyBytesObject *)self->co_code;
+    PyObject *co_consts = self->co_consts;
+    PyObject *co_names = self->co_names;
+    PyObject *co_varnames = self->co_varnames;
+    PyObject *co_freevars = self->co_freevars;
+    PyObject *co_cellvars = self->co_cellvars;
+    PyObject *co_filename = self->co_filename;
+    PyObject *co_name = self->co_name;
+    PyBytesObject *co_lnotab = (PyBytesObject *)self->co_lnotab;
+
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 0, 0, argsbuf);
+    if (!args) {
+        goto exit;
+    }
+    if (!noptargs) {
+        goto skip_optional_kwonly;
+    }
+    if (args[0]) {
+        if (PyFloat_Check(args[0])) {
+            PyErr_SetString(PyExc_TypeError,
+                            "integer argument expected, got float" );
+            goto exit;
+        }
+        co_argcount = _PyLong_AsInt(args[0]);
+        if (co_argcount == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[1]) {
+        if (PyFloat_Check(args[1])) {
+            PyErr_SetString(PyExc_TypeError,
+                            "integer argument expected, got float" );
+            goto exit;
+        }
+        co_posonlyargcount = _PyLong_AsInt(args[1]);
+        if (co_posonlyargcount == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[2]) {
+        if (PyFloat_Check(args[2])) {
+            PyErr_SetString(PyExc_TypeError,
+                            "integer argument expected, got float" );
+            goto exit;
+        }
+        co_kwonlyargcount = _PyLong_AsInt(args[2]);
+        if (co_kwonlyargcount == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[3]) {
+        if (PyFloat_Check(args[3])) {
+            PyErr_SetString(PyExc_TypeError,
+                            "integer argument expected, got float" );
+            goto exit;
+        }
+        co_nlocals = _PyLong_AsInt(args[3]);
+        if (co_nlocals == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[4]) {
+        if (PyFloat_Check(args[4])) {
+            PyErr_SetString(PyExc_TypeError,
+                            "integer argument expected, got float" );
+            goto exit;
+        }
+        co_stacksize = _PyLong_AsInt(args[4]);
+        if (co_stacksize == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[5]) {
+        if (PyFloat_Check(args[5])) {
+            PyErr_SetString(PyExc_TypeError,
+                            "integer argument expected, got float" );
+            goto exit;
+        }
+        co_flags = _PyLong_AsInt(args[5]);
+        if (co_flags == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[6]) {
+        if (PyFloat_Check(args[6])) {
+            PyErr_SetString(PyExc_TypeError,
+                            "integer argument expected, got float" );
+            goto exit;
+        }
+        co_firstlineno = _PyLong_AsInt(args[6]);
+        if (co_firstlineno == -1 && PyErr_Occurred()) {
+            goto exit;
+        }
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[7]) {
+        if (!PyBytes_Check(args[7])) {
+            _PyArg_BadArgument("replace", 8, "bytes", args[7]);
+            goto exit;
+        }
+        co_code = (PyBytesObject *)args[7];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[8]) {
+        if (!PyTuple_Check(args[8])) {
+            _PyArg_BadArgument("replace", 9, "tuple", args[8]);
+            goto exit;
+        }
+        co_consts = args[8];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[9]) {
+        if (!PyTuple_Check(args[9])) {
+            _PyArg_BadArgument("replace", 10, "tuple", args[9]);
+            goto exit;
+        }
+        co_names = args[9];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[10]) {
+        if (!PyTuple_Check(args[10])) {
+            _PyArg_BadArgument("replace", 11, "tuple", args[10]);
+            goto exit;
+        }
+        co_varnames = args[10];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[11]) {
+        if (!PyTuple_Check(args[11])) {
+            _PyArg_BadArgument("replace", 12, "tuple", args[11]);
+            goto exit;
+        }
+        co_freevars = args[11];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[12]) {
+        if (!PyTuple_Check(args[12])) {
+            _PyArg_BadArgument("replace", 13, "tuple", args[12]);
+            goto exit;
+        }
+        co_cellvars = args[12];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[13]) {
+        if (!PyUnicode_Check(args[13])) {
+            _PyArg_BadArgument("replace", 14, "str", args[13]);
+            goto exit;
+        }
+        if (PyUnicode_READY(args[13]) == -1) {
+            goto exit;
+        }
+        co_filename = args[13];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (args[14]) {
+        if (!PyUnicode_Check(args[14])) {
+            _PyArg_BadArgument("replace", 15, "str", args[14]);
+            goto exit;
+        }
+        if (PyUnicode_READY(args[14]) == -1) {
+            goto exit;
+        }
+        co_name = args[14];
+        if (!--noptargs) {
+            goto skip_optional_kwonly;
+        }
+    }
+    if (!PyBytes_Check(args[15])) {
+        _PyArg_BadArgument("replace", 16, "bytes", args[15]);
+        goto exit;
+    }
+    co_lnotab = (PyBytesObject *)args[15];
+skip_optional_kwonly:
+    return_value = code_replace_impl(self, co_argcount, co_posonlyargcount, co_kwonlyargcount, co_nlocals, co_stacksize, co_flags, co_firstlineno, co_code, co_consts, co_names, co_varnames, co_freevars, co_cellvars, co_filename, co_name, co_lnotab);
+
+exit:
+    return return_value;
+}
+/*[clinic end generated code: output=624ab6f2ea8f0ea4 input=a9049054013a1b77]*/
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index f4e48a9..1e76f26 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -5,6 +5,7 @@
 #include "structmember.h"
 #include "pycore_pystate.h"
 #include "pycore_tupleobject.h"
+#include "clinic/codeobject.c.h"
 
 /* Holder for co_extra information */
 typedef struct {
@@ -12,6 +13,11 @@ typedef struct {
     void *ce_extras[1];
 } _PyCodeObjectExtra;
 
+/*[clinic input]
+class code "PyCodeObject *" "&PyCode_Type"
+[clinic start generated code]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=78aa5d576683bb4b]*/
+
 /* all_name_chars(s): true iff s matches [a-zA-Z0-9_]* */
 static int
 all_name_chars(PyObject *o)
@@ -109,7 +115,8 @@ PyCode_New(int argcount, int posonlyargcount, int kwonlyargcount,
 
     /* Check argument types */
     if (argcount < 0 || posonlyargcount < 0 || kwonlyargcount < 0 ||
-        nlocals < 0 || code == NULL || !PyBytes_Check(code) ||
+        nlocals < 0 || stacksize < 0 || flags < 0 ||
+        code == NULL || !PyBytes_Check(code) ||
         consts == NULL || !PyTuple_Check(consts) ||
         names == NULL || !PyTuple_Check(names) ||
         varnames == NULL || !PyTuple_Check(varnames) ||
@@ -122,9 +129,13 @@ PyCode_New(int argcount, int posonlyargcount, int kwonlyargcount,
         return NULL;
     }
 
-    /* Ensure that the filename is a ready Unicode string */
-    if (PyUnicode_READY(filename) < 0)
+    /* Ensure that strings are ready Unicode string */
+    if (PyUnicode_READY(name) < 0) {
+        return NULL;
+    }
+    if (PyUnicode_READY(filename) < 0) {
         return NULL;
+    }
 
     intern_strings(names);
     intern_strings(varnames);
@@ -482,7 +493,7 @@ code_dealloc(PyCodeObject *co)
 }
 
 static PyObject *
-code_sizeof(PyCodeObject *co, void *unused)
+code_sizeof(PyCodeObject *co, PyObject *Py_UNUSED(args))
 {
     Py_ssize_t res = _PyObject_SIZE(Py_TYPE(co));
     _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra*) co->co_extra;
@@ -497,6 +508,65 @@ code_sizeof(PyCodeObject *co, void *unused)
     return PyLong_FromSsize_t(res);
 }
 
+/*[clinic input]
+code.replace
+
+    *
+    co_argcount: int(c_default="self->co_argcount") = -1
+    co_posonlyargcount: int(c_default="self->co_posonlyargcount") = -1
+    co_kwonlyargcount: int(c_default="self->co_kwonlyargcount") = -1
+    co_nlocals: int(c_default="self->co_nlocals") = -1
+    co_stacksize: int(c_default="self->co_stacksize") = -1
+    co_flags: int(c_default="self->co_flags") = -1
+    co_firstlineno: int(c_default="self->co_firstlineno") = -1
+    co_code: PyBytesObject(c_default="(PyBytesObject *)self->co_code") = None
+    co_consts: object(subclass_of="&PyTuple_Type", c_default="self->co_consts") = None
+    co_names: object(subclass_of="&PyTuple_Type", c_default="self->co_names") = None
+    co_varnames: object(subclass_of="&PyTuple_Type", c_default="self->co_varnames") = None
+    co_freevars: object(subclass_of="&PyTuple_Type", c_default="self->co_freevars") = None
+    co_cellvars: object(subclass_of="&PyTuple_Type", c_default="self->co_cellvars") = None
+    co_filename: unicode(c_default="self->co_filename") = None
+    co_name: unicode(c_default="self->co_name") = None
+    co_lnotab: PyBytesObject(c_default="(PyBytesObject *)self->co_lnotab") = None
+
+Return a new code object with new specified fields.
+[clinic start generated code]*/
+
+static PyObject *
+code_replace_impl(PyCodeObject *self, int co_argcount,
+                  int co_posonlyargcount, int co_kwonlyargcount,
+                  int co_nlocals, int co_stacksize, int co_flags,
+                  int co_firstlineno, PyBytesObject *co_code,
+                  PyObject *co_consts, PyObject *co_names,
+                  PyObject *co_varnames, PyObject *co_freevars,
+                  PyObject *co_cellvars, PyObject *co_filename,
+                  PyObject *co_name, PyBytesObject *co_lnotab)
+/*[clinic end generated code: output=25c8e303913bcace input=77189e46579ec426]*/
+{
+#define CHECK_INT_ARG(ARG) \
+        if (ARG < 0) { \
+            PyErr_SetString(PyExc_ValueError, \
+                            #ARG " must be a positive integer"); \
+            return NULL; \
+        }
+
+    CHECK_INT_ARG(co_argcount);
+    CHECK_INT_ARG(co_posonlyargcount);
+    CHECK_INT_ARG(co_kwonlyargcount);
+    CHECK_INT_ARG(co_nlocals);
+    CHECK_INT_ARG(co_stacksize);
+    CHECK_INT_ARG(co_flags);
+    CHECK_INT_ARG(co_firstlineno);
+
+#undef CHECK_INT_ARG
+
+    return (PyObject *)PyCode_New(
+        co_argcount, co_posonlyargcount, co_kwonlyargcount, co_nlocals,
+        co_stacksize, co_flags, (PyObject*)co_code, co_consts, co_names,
+        co_varnames, co_freevars, co_cellvars, co_filename, co_name,
+        co_firstlineno, (PyObject*)co_lnotab);
+}
+
 static PyObject *
 code_repr(PyCodeObject *co)
 {
@@ -751,6 +821,7 @@ code_hash(PyCodeObject *co)
 
 static struct PyMethodDef code_methods[] = {
     {"__sizeof__", (PyCFunction)code_sizeof, METH_NOARGS},
+    CODE_REPLACE_METHODDEF
     {NULL, NULL}                /* sentinel */
 };