Fix test_code: compiler and code object improvements

youknowone · youknowone · commit c57f4decd590 · 2026-03-17T10:54:12.000+09:00
- Add CO_NESTED flag (0x10) for nested function scopes
- Emit LOAD_SMALL_INT for integers 0..=255 instead of LOAD_CONST
- Eliminate dead constant expression statements (no side effects)
- Ensure None in co_consts for functions with no other constants
- Add code.__replace__() for copy.replace() support
- Mark test_co_lnotab and test_invalid_bytecode as expectedFailure
diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
@@ -412,8 +412,6 @@ def func():
         with self.assertRaises(ValueError):
             co.replace(co_nlocals=co.co_nlocals + 1)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_shrinking_localsplus(self):
         # Check that PyCode_NewWithPosOnlyArgs resizes both
         # localsplusnames and localspluskinds, if an argument is a cell.
@@ -429,7 +427,7 @@ def func():
         new_code = code = func.__code__.replace(co_linetable=b'')
         self.assertEqual(list(new_code.co_lines()), [])
 
-    # TODO: RUSTPYTHON
+    # TODO: RUSTPYTHON; co_lnotab intentionally not implemented (deprecated since 3.12)
     @unittest.expectedFailure
     def test_co_lnotab_is_deprecated(self):  # TODO: remove in 3.14
         def func():
@@ -493,6 +491,8 @@ def spam9():
                 res = _testinternalcapi.code_returns_only_none(func.__code__)
                 self.assertFalse(res)
 
+    # TODO: RUSTPYTHON; replace() rejects invalid bytecodes for safety
+    @unittest.expectedFailure
     def test_invalid_bytecode(self):
         def foo():
             pass
diff --git a/Lib/test/test_codeop.py b/Lib/test/test_codeop.py
@@ -30,7 +30,6 @@ def assertInvalid(self, str, symbol='single', is_syntax=1):
         except OverflowError:
             self.assertTrue(not is_syntax)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: <code object <module> at 0xc99532080 file "<input>", line 1> != <code object <module> at 0xc99532f80 file "<input>", line 1>
     def test_valid(self):
         av = self.assertValid
 
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -1131,7 +1131,6 @@ def test_bug_46724(self):
         # Test that negative operargs are handled properly
         self.do_disassembly_test(bug46724, dis_bug46724)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_kw_names(self):
         # Test that value is displayed for keyword argument names:
         self.do_disassembly_test(wrap_func_w_kwargs, dis_kw_names)
@@ -1179,7 +1178,6 @@ def test_disassemble_str(self):
         self.do_disassembly_test(fn_with_annotate_str, dis_fn_with_annotate_str)
         self.do_disassembly_test(compound_stmt_str, dis_compound_stmt_str)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_disassemble_bytes(self):
         self.do_disassembly_test(_f.__code__.co_code, dis_f_co_code)
 
diff --git a/Lib/test/test_marshal.py b/Lib/test/test_marshal.py
@@ -123,7 +123,6 @@ def test_exceptions(self):
         self.assertEqual(StopIteration, new)
 
 class CodeTestCase(unittest.TestCase):
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_code(self):
         co = ExceptionTestCase.test_exceptions.__code__
         new = marshal.loads(marshal.dumps(co))
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
@@ -1081,6 +1081,14 @@ impl Compiler {
             ),
         };
 
+        // Set CO_NESTED for scopes defined inside another function/class/etc.
+        // (i.e., not at module level)
+        let flags = if self.code_stack.len() > 1 {
+            flags | bytecode::CodeFlags::NESTED
+        } else {
+            flags
+        };
+
         // Get private name from parent scope
         let private = if !self.code_stack.is_empty() {
             self.code_stack.last().unwrap().private.clone()
@@ -1202,7 +1210,8 @@ impl Compiler {
         // enter_scope sets default values based on scope_type, but push_output
         // allows callers to specify exact values
         if let Some(info) = self.code_stack.last_mut() {
-            info.flags = flags;
+            // Preserve NESTED flag set by enter_scope
+            info.flags = flags | (info.flags & bytecode::CodeFlags::NESTED);
             info.metadata.argcount = arg_count;
             info.metadata.posonlyargcount = posonlyarg_count;
             info.metadata.kwonlyargcount = kwonlyarg_count;
@@ -2179,18 +2188,26 @@ impl Compiler {
                 }
             }
             ast::Stmt::Expr(ast::StmtExpr { value, .. }) => {
-                self.compile_expression(value)?;
+                // Optimize away constant expressions with no side effects.
+                // In interactive mode, always compile (to print the result).
+                let dominated_by_interactive =
+                    self.interactive && !self.ctx.in_func() && !self.ctx.in_class;
+                if !dominated_by_interactive && Self::is_const_expression(value) {
+                    // Skip compilation entirely - the expression has no side effects
+                } else {
+                    self.compile_expression(value)?;
 
-                if self.interactive && !self.ctx.in_func() && !self.ctx.in_class {
-                    emit!(
-                        self,
-                        Instruction::CallIntrinsic1 {
-                            func: bytecode::IntrinsicFunction1::Print
-                        }
-                    );
-                }
+                    if dominated_by_interactive {
+                        emit!(
+                            self,
+                            Instruction::CallIntrinsic1 {
+                                func: bytecode::IntrinsicFunction1::Print
+                            }
+                        );
+                    }
 
-                emit!(self, Instruction::PopTop);
+                    emit!(self, Instruction::PopTop);
+                }
             }
             ast::Stmt::Global(_) | ast::Stmt::Nonlocal(_) => {
                 // Handled during symbol table construction.
@@ -3748,19 +3765,23 @@ impl Compiler {
                 });
             self.current_code_info().flags |= bytecode::CodeFlags::HAS_DOCSTRING;
         }
-        // If no docstring, don't add None to co_consts
-        // Note: RETURN_GENERATOR + POP_TOP for async functions is emitted in enter_scope()
-
         // Compile body statements
         self.compile_statements(body)?;
 
-        // Emit None at end if needed
+        // Emit implicit `return None` if the body doesn't end with return.
+        // Also ensure None is in co_consts even when not emitting return
+        // (matching CPython: functions without explicit constants always
+        // have None in co_consts).
         match body.last() {
             Some(ast::Stmt::Return(_)) => {}
             _ => {
                 self.emit_return_const(ConstantData::None);
             }
         }
+        // Functions with no other constants should still have None in co_consts
+        if self.current_code_info().metadata.consts.is_empty() {
+            self.arg_constant(ConstantData::None);
+        }
 
         // Exit scope and create function object
         let code = self.exit_scope();
@@ -6882,6 +6903,19 @@ impl Compiler {
         Ok(send_block)
     }
 
+    /// Returns true if the expression is a constant with no side effects.
+    fn is_const_expression(expr: &ast::Expr) -> bool {
+        matches!(
+            expr,
+            ast::Expr::StringLiteral(_)
+                | ast::Expr::BytesLiteral(_)
+                | ast::Expr::NumberLiteral(_)
+                | ast::Expr::BooleanLiteral(_)
+                | ast::Expr::NoneLiteral(_)
+                | ast::Expr::EllipsisLiteral(_)
+        )
+    }
+
     fn compile_expression(&mut self, expression: &ast::Expr) -> CompileResult<()> {
         trace!("Compiling {expression:?}");
         let range = expression.range();
diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs
@@ -190,13 +190,7 @@ impl CodeInfo {
     ) -> crate::InternalResult<CodeObject> {
         // Always fold tuple constants
         self.fold_tuple_constants();
-        // Python only applies LOAD_SMALL_INT conversion to module-level code
-        // (not inside functions). Module code lacks OPTIMIZED flag.
-        // Note: RustPython incorrectly sets NEWLOCALS on modules, so only check OPTIMIZED
-        let is_module_level = !self.flags.contains(CodeFlags::OPTIMIZED);
-        if is_module_level {
-            self.convert_to_load_small_int();
-        }
+        self.convert_to_load_small_int();
         self.remove_unused_consts();
         self.remove_nops();
 
@@ -786,8 +780,8 @@ impl CodeInfo {
                     continue;
                 };
 
-                // Check if it's in small int range: -5 to 256 (_PY_IS_SMALL_INT)
-                if let Some(small) = value.to_i32().filter(|v| (-5..=256).contains(v)) {
+                // LOAD_SMALL_INT oparg is unsigned, so only 0..=255 can be encoded
+                if let Some(small) = value.to_i32().filter(|v| (0..=255).contains(v)) {
                     // Convert LOAD_CONST to LOAD_SMALL_INT
                     instr.instr = Instruction::LoadSmallInt { i: Arg::marker() }.into();
                     // The arg is the i32 value stored as u32 (two's complement)
diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs
@@ -371,6 +371,7 @@ bitflags! {
         const NEWLOCALS = 0x0002;
         const VARARGS = 0x0004;
         const VARKEYWORDS = 0x0008;
+        const NESTED = 0x0010;
         const GENERATOR = 0x0020;
         const COROUTINE = 0x0080;
         const ITERABLE_COROUTINE = 0x0100;
diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs
@@ -664,7 +664,10 @@ impl Constructor for PyCode {
     }
 }
 
-#[pyclass(with(Representable, Constructor, Comparable, Hashable), flags(HAS_WEAKREF))]
+#[pyclass(
+    with(Representable, Constructor, Comparable, Hashable),
+    flags(HAS_WEAKREF)
+)]
 impl PyCode {
     #[pygetset]
     const fn co_posonlyargcount(&self) -> usize {
@@ -790,6 +793,7 @@ impl PyCode {
         vm.ctx.new_bytes(self.code.exceptiontable.to_vec())
     }
 
+    // spell-checker: ignore lnotab
     // co_lnotab is intentionally not implemented.
     // It was deprecated since 3.12 and scheduled for removal in 3.14.
     // Use co_lines() or co_linetable instead.
@@ -1065,6 +1069,11 @@ impl PyCode {
         vm.call_method(list.as_object(), "__iter__", ())
     }
 
+    #[pymethod]
+    pub fn __replace__(&self, args: ReplaceArgs, vm: &VirtualMachine) -> PyResult<Self> {
+        self.replace(args, vm)
+    }
+
     #[pymethod]
     pub fn replace(&self, args: ReplaceArgs, vm: &VirtualMachine) -> PyResult<Self> {
         let ReplaceArgs {
