Fix test_code: compiler and code object improvements

youknowone · youknowone · commit 75bece0c2a51 · 2026-03-16T12:02:47.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
@@ -429,7 +429,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 +493,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/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();
@@ -8090,6 +8124,13 @@ impl Compiler {
     }
 
     fn emit_load_const(&mut self, constant: ConstantData) {
+        // Use LOAD_SMALL_INT for integers 0..=255
+        if let ConstantData::Integer { ref value } = constant {
+            if let Some(small) = value.try_into().ok().filter(|v: &u32| *v <= 255) {
+                emit!(self, Instruction::LoadSmallInt { i: small });
+                return;
+            }
+        }
         let idx = self.arg_constant(constant);
         self.emit_arg(idx, |consti| Instruction::LoadConst { consti })
     }
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 {
@@ -1065,6 +1068,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 {
