Fix a SystemError in code.replace() (#27771)

gvanrossum · web-flow · commit 62bd97303eb6 · 2021-08-16T11:34:23.000-07:00
While the comment said 'We don't bother resizing localspluskinds',
this would cause .replace() to crash when it happened.
(Also types.CodeType(), but testing that is tedious, and this tests all
code paths.)
diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
@@ -320,6 +320,15 @@ def func():
         with self.assertRaises(ValueError):
             co.replace(co_nlocals=co.co_nlocals + 1)
 
+    def test_shrinking_localsplus(self):
+        # Check that PyCode_NewWithPosOnlyArgs resizes both
+        # localsplusnames and localspluskinds, if an argument is a cell.
+        def func(arg):
+            return lambda: arg
+        code = func.__code__
+        newcode = code.replace(co_name="func")  # Should not raise SystemError
+        self.assertEqual(code, newcode)
+
     def test_empty_linetable(self):
         def func():
             pass
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
@@ -471,9 +471,11 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
                                localsplusnames, localspluskinds);
     }
     // If any cells were args then nlocalsplus will have shrunk.
-    // We don't bother resizing localspluskinds.
-    if (_PyTuple_Resize(&localsplusnames, nlocalsplus) < 0) {
-        goto error;
+    if (nlocalsplus != PyTuple_GET_SIZE(localsplusnames)) {
+        if (_PyTuple_Resize(&localsplusnames, nlocalsplus) < 0
+                || _PyBytes_Resize(&localspluskinds, nlocalsplus) < 0) {
+            goto error;
+        }
     }
 
     struct _PyCodeConstructor con = {

Original file line number	Diff line number	Diff line change
`@@ -471,9 +471,11 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,`
`471`	`471`	`localsplusnames, localspluskinds);`
`472`	`472`	`}`
`473`	`473`	`// If any cells were args then nlocalsplus will have shrunk.`
`474`		`- // We don't bother resizing localspluskinds.`
`475`		`- if (_PyTuple_Resize(&localsplusnames, nlocalsplus) < 0) {`
`476`		`- goto error;`
	`474`	`+ if (nlocalsplus != PyTuple_GET_SIZE(localsplusnames)) {`
	`475`	`+ if (_PyTuple_Resize(&localsplusnames, nlocalsplus) < 0`
	`476`	`+ \|\| _PyBytes_Resize(&localspluskinds, nlocalsplus) < 0) {`
	`477`	`+ goto error;`
	`478`	`+ }`
`477`	`479`	`}`
`478`	`480`
`479`	`481`	`struct _PyCodeConstructor con = {`