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

4 files changed, 64 insertions, 28 deletions

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