Merge pull request #703 from RustPython/add_class_cell

windelbouwman · web-flow · commit 4c485202d15a · 2019-03-20T13:56:03.000+01:00
Add __class__ cell to method scopes.
diff --git a/tests/snippets/class.py b/tests/snippets/class.py
@@ -22,15 +22,18 @@ def __init__(self, x):
         self.x = x
 
     def get_x(self):
+        assert __class__ is Bar
         return self.x
 
     @classmethod
     def fubar(cls, x):
+        assert __class__ is cls
         assert cls is Bar
         assert x == 2
 
     @staticmethod
     def kungfu(x):
+        assert __class__ is Bar
         assert x == 3
 
 
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -20,7 +20,7 @@ use crate::obj::objtype;
 use crate::frame::Scope;
 use crate::function::{Args, OptionalArg, PyFuncArgs};
 use crate::pyobject::{
-    AttributeProtocol, IdProtocol, PyContext, PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, DictProtocol, IdProtocol, PyContext, PyObjectRef, PyResult, TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 
@@ -842,7 +842,10 @@ pub fn builtin_build_class_(vm: &mut VirtualMachine, mut args: PyFuncArgs) -> Py
     let prepare = vm.get_attribute(metaclass.clone(), prepare_name)?;
     let namespace = vm.invoke(prepare, vec![name_arg.clone(), bases.clone()])?;
 
-    vm.invoke_with_locals(function, namespace.clone())?;
+    let cells = vm.new_dict();
 
-    vm.call_method(&metaclass, "__call__", vec![name_arg, bases, namespace])
+    vm.invoke_with_locals(function, cells.clone(), namespace.clone())?;
+    let class = vm.call_method(&metaclass, "__call__", vec![name_arg, bases, namespace])?;
+    cells.set_item(&vm.ctx, "__class__", class.clone());
+    Ok(class)
 }
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -359,14 +359,21 @@ impl VirtualMachine {
         }
     }
 
-    pub fn invoke_with_locals(&mut self, function: PyObjectRef, locals: PyObjectRef) -> PyResult {
+    pub fn invoke_with_locals(
+        &mut self,
+        function: PyObjectRef,
+        cells: PyObjectRef,
+        locals: PyObjectRef,
+    ) -> PyResult {
         if let Some(PyFunction {
             code,
             scope,
             defaults: _,
         }) = &function.payload()
         {
-            let scope = scope.child_scope_with_locals(locals);
+            let scope = scope
+                .child_scope_with_locals(cells)
+                .child_scope_with_locals(locals);
             let frame = self.ctx.new_frame(code.clone(), scope);
             return self.run_frame_full(frame);
         }

Original file line number	Diff line number	Diff line change
`@@ -359,14 +359,21 @@ impl VirtualMachine {`
`359`	`359`	`}`
`360`	`360`	`}`
`361`	`361`
`362`		`- pub fn invoke_with_locals(&mut self, function: PyObjectRef, locals: PyObjectRef) -> PyResult {`
	`362`	`+ pub fn invoke_with_locals(`
	`363`	`+ &mut self,`
	`364`	`+ function: PyObjectRef,`
	`365`	`+ cells: PyObjectRef,`
	`366`	`+ locals: PyObjectRef,`
	`367`	`+ ) -> PyResult {`
`363`	`368`	`if let Some(PyFunction {`
`364`	`369`	`code,`
`365`	`370`	`scope,`
`366`	`371`	`defaults: _,`
`367`	`372`	`}) = &function.payload()`
`368`	`373`	`{`
`369`		`- let scope = scope.child_scope_with_locals(locals);`
	`374`	`+ let scope = scope`
	`375`	`+ .child_scope_with_locals(cells)`
	`376`	`+ .child_scope_with_locals(locals);`
`370`	`377`	`let frame = self.ctx.new_frame(code.clone(), scope);`
`371`	`378`	`return self.run_frame_full(frame);`
`372`	`379`	`}`