RustPython · windelbouwman · Mar 20, 2019 · Mar 20, 2019
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);
         }