Only overwrite the locals in __build_class__ if it's a class

Author: coolreader18

bytecode/src/bytecode.rs

@@ -53,6 +53,7 @@ bitflags! {
         const HAS_DEFAULTS = 0x01;
         const HAS_KW_ONLY_DEFAULTS = 0x02;
         const HAS_ANNOTATIONS = 0x04;
+        const NO_NEW_LOCALS = 0x08;
     }
 }
 

compiler/src/compile.rs

@@ -955,12 +955,21 @@ impl<O: OutputStream> Compiler<O> {
 
         self.emit(Instruction::LoadName {
             name: "__name__".to_string(),
-            scope: bytecode::NameScope::Free,
+            scope: bytecode::NameScope::Global,
         });
         self.emit(Instruction::StoreName {
             name: "__module__".to_string(),
             scope: bytecode::NameScope::Free,
         });
+        self.emit(Instruction::LoadConst {
+            value: bytecode::Constant::String {
+                value: qualified_name.clone(),
+            },
+        });
+        self.emit(Instruction::StoreName {
+            name: "__qualname__".to_string(),
+            scope: bytecode::NameScope::Free,
+        });
         self.compile_statements(new_body)?;
         self.emit(Instruction::LoadConst {
             value: bytecode::Constant::None,
@@ -983,7 +992,7 @@ impl<O: OutputStream> Compiler<O> {
 
         // Turn code object into function object:
         self.emit(Instruction::MakeFunction {
-            flags: bytecode::FunctionOpArg::empty(),
+            flags: bytecode::FunctionOpArg::NO_NEW_LOCALS,
         });
 
         self.emit(Instruction::LoadConst {

compiler/src/symboltable.rs

@@ -387,6 +387,8 @@ impl SymbolTableBuilder {
                 decorator_list,
             } => {
                 self.enter_scope(name, SymbolTableType::Class, statement.location.row());
+                self.register_name("__module__", SymbolUsage::Assigned)?;
+                self.register_name("__qualname__", SymbolUsage::Assigned)?;
                 self.scan_statements(body)?;
                 self.leave_scope();
                 self.scan_expressions(bases, &ExpressionContext::Load)?;

vm/src/builtins.rs

@@ -14,6 +14,7 @@ use crate::obj::objbool::{self, IntoPyBool};
 use crate::obj::objbytes::PyBytesRef;
 use crate::obj::objcode::PyCodeRef;
 use crate::obj::objdict::PyDictRef;
+use crate::obj::objfunction::PyFunctionRef;
 use crate::obj::objint::{self, PyIntRef};
 use crate::obj::objiter;
 use crate::obj::objstr::{PyString, PyStringRef};
@@ -832,6 +833,7 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
         "exit" => ctx.new_rustfunc(builtin_exit),
         "quit" => ctx.new_rustfunc(builtin_exit),
         "__import__" => ctx.new_rustfunc(builtin_import),
+        "__build_class__" => ctx.new_rustfunc(builtin_build_class_),
 
         // Constants
         "NotImplemented" => ctx.not_implemented(),
@@ -888,7 +890,7 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
 }
 
 pub fn builtin_build_class_(
-    function: PyObjectRef,
+    function: PyFunctionRef,
     qualified_name: PyStringRef,
     bases: Args<PyClassRef>,
     mut kwargs: KwArgs,
@@ -925,10 +927,12 @@ pub fn builtin_build_class_(
 
     let cells = vm.ctx.new_dict();
 
-    vm.invoke_with_locals(&function, cells.clone(), namespace.clone())?;
+    let scope = function
+        .scope
+        .new_child_scope_with_locals(cells.clone())
+        .new_child_scope_with_locals(namespace.clone());
 
-    namespace.set_item("__name__", name_obj.clone(), vm)?;
-    namespace.set_item("__qualname__", qualified_name.into_object(), vm)?;
+    vm.invoke_python_function_with_scope(&function, vec![].into(), &scope)?;
 
     let class = vm.call_method(
         metaclass.as_object(),

vm/src/frame.rs

@@ -1,7 +1,6 @@
 use std::cell::RefCell;
 use std::fmt;
 
-use crate::builtins;
 use crate::bytecode;
 use crate::function::PyFuncArgs;
 use crate::obj::objbool;
@@ -460,7 +459,7 @@ impl Frame {
                 Ok(None)
             }
             bytecode::Instruction::LoadBuildClass => {
-                self.push_value(vm.ctx.new_rustfunc(builtins::builtin_build_class_));
+                self.push_value(vm.get_attribute(vm.builtins.clone(), "__build_class__")?);
                 Ok(None)
             }
             bytecode::Instruction::UnpackSequence { size } => {
@@ -1054,9 +1053,13 @@ impl Frame {
         // pop argc arguments
         // argument: name, args, globals
         let scope = self.scope.clone();
-        let func_obj = vm
-            .ctx
-            .new_function(code_obj, scope, defaults, kw_only_defaults);
+        let func_obj = vm.ctx.new_function(
+            code_obj,
+            scope,
+            defaults,
+            kw_only_defaults,
+            !flags.contains(bytecode::FunctionOpArg::NO_NEW_LOCALS),
+        );
 
         let name = qualified_name.as_str().split('.').next_back().unwrap();
         vm.set_attr(&func_obj, "__name__", vm.new_str(name.to_string()))?;

vm/src/obj/objfunction.rs

@@ -16,6 +16,7 @@ pub struct PyFunction {
     pub scope: Scope,
     pub defaults: Option<PyTupleRef>,
     pub kw_only_defaults: Option<PyDictRef>,
+    pub new_locals: bool,
 }
 
 impl PyFunction {
@@ -24,12 +25,14 @@ impl PyFunction {
         scope: Scope,
         defaults: Option<PyTupleRef>,
         kw_only_defaults: Option<PyDictRef>,
+        new_locals: bool,
     ) -> Self {
         PyFunction {
             code,
             scope,
             defaults,
             kw_only_defaults,
+            new_locals,
         }
     }
 }

vm/src/pyobject.rs

@@ -478,9 +478,10 @@ impl PyContext {
         scope: Scope,
         defaults: Option<PyTupleRef>,
         kw_only_defaults: Option<PyDictRef>,
+        new_locals: bool,
     ) -> PyObjectRef {
         PyObject::new(
-            PyFunction::new(code_obj, scope, defaults, kw_only_defaults),
+            PyFunction::new(code_obj, scope, defaults, kw_only_defaults, new_locals),
             self.function_type(),
             Some(self.new_dict()),
         )

vm/src/vm.rs

@@ -570,15 +570,9 @@ impl VirtualMachine {
     fn _invoke(&self, func_ref: &PyObjectRef, args: PyFuncArgs) -> PyResult {
         vm_trace!("Invoke: {:?} {:?}", func_ref, args);
 
-        if let Some(PyFunction {
-            ref code,
-            ref scope,
-            ref defaults,
-            ref kw_only_defaults,
-        }) = func_ref.payload()
-        {
+        if let Some(py_func) = func_ref.payload() {
             self.trace_event(TraceEvent::Call)?;
-            let res = self.invoke_python_function(code, scope, defaults, kw_only_defaults, args);
+            let res = self.invoke_python_function(py_func, args);
             self.trace_event(TraceEvent::Return)?;
             res
         } else if let Some(PyMethod {
@@ -634,21 +628,30 @@ impl VirtualMachine {
         Ok(())
     }
 
-    fn invoke_python_function(
+    pub fn invoke_python_function(&self, func: &PyFunction, func_args: PyFuncArgs) -> PyResult {
+        self.invoke_python_function_with_scope(func, func_args, &func.scope)
+    }
+
+    pub fn invoke_python_function_with_scope(
         &self,
-        code: &PyCodeRef,
-        scope: &Scope,
-        defaults: &Option<PyTupleRef>,
-        kw_only_defaults: &Option<PyDictRef>,
+        func: &PyFunction,
         func_args: PyFuncArgs,
+        scope: &Scope,
     ) -> PyResult {
-        let scope = scope.new_child_scope(&self.ctx);
+        let code = &func.code;
+
+        let scope = if func.new_locals {
+            scope.new_child_scope(&self.ctx)
+        } else {
+            scope.clone()
+        };
+
         self.fill_locals_from_args(
             &code.code,
             &scope.get_locals(),
             func_args,
-            defaults,
-            kw_only_defaults,
+            &func.defaults,
+            &func.kw_only_defaults,
         )?;
 
         // Construct frame:
@@ -662,25 +665,6 @@ impl VirtualMachine {
         }
     }
 
-    pub fn invoke_with_locals(
-        &self,
-        function: &PyObjectRef,
-        cells: PyDictRef,
-        locals: PyDictRef,
-    ) -> PyResult {
-        if let Some(PyFunction { code, scope, .. }) = &function.payload() {
-            let scope = scope
-                .new_child_scope_with_locals(cells)
-                .new_child_scope_with_locals(locals);
-            let frame = Frame::new(code.clone(), scope).into_ref(self);
-            return self.run_frame_full(frame);
-        }
-        panic!(
-            "invoke_with_locals: expected python function, got: {:?}",
-            *function
-        );
-    }
-
     fn fill_locals_from_args(
         &self,
         code_object: &bytecode::CodeObject,