From 1958e470221393167690f270d2969b237c225a22 Mon Sep 17 00:00:00 2001
From: Adam Kelly <adam@cthulahoops.org>
Date: Wed, 20 Mar 2019 09:20:28 +0000
Subject: [PATCH] Add __class__ cell to method scopes.

---
 tests/snippets/class.py |  3 +++
 vm/src/builtins.rs      |  9 ++++++---
 vm/src/vm.rs            | 11 +++++++++--
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/tests/snippets/class.py b/tests/snippets/class.py
index 42144d720bb..84ed872460a 100644
--- 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
index dccaeee841b..432d661a073 100644
--- 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
index 593c15ea932..47cc162a8f0 100644
--- 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);
         }