From e513bf11fee3836d08fe23a04799130f321a4739 Mon Sep 17 00:00:00 2001
From: YeonWoo <83155697+oow214@users.noreply.github.com>
Date: Sat, 13 Aug 2022 17:20:13 +0900
Subject: [PATCH] Fix metaclass check classcell (#3986)

---
 Lib/test/test_super.py    |  6 ------
 vm/src/builtins/type.rs   | 14 +++++++++++++-
 vm/src/stdlib/builtins.rs | 14 +++++++++++++-
 vm/src/vm/context.rs      |  1 +
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/Lib/test/test_super.py b/Lib/test/test_super.py
index b5f142985..ca8dc6e86 100644
--- a/Lib/test/test_super.py
+++ b/Lib/test/test_super.py
@@ -143,8 +143,6 @@ class TestSuper(unittest.TestCase):
                 return __class__
         self.assertIs(X.f(), X)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test___class___new(self):
         # See issue #23722
         # Ensure zero-arg super() works as soon as type.__new__() is completed
@@ -263,8 +261,6 @@ class TestSuper(unittest.TestCase):
                 def f(self):
                     return __class__
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test___classcell___overwrite(self):
         # See issue #23722
         # Overwriting __classcell__ with nonsense is explicitly prohibited
@@ -279,8 +275,6 @@ class TestSuper(unittest.TestCase):
                     class A(metaclass=Meta, cell=bad_cell):
                         pass
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test___classcell___wrong_cell(self):
         # See issue #23722
         # Pointing the cell reference at the wrong class is also prohibited
diff --git a/vm/src/builtins/type.rs b/vm/src/builtins/type.rs
index f1f0300f1..81357e0f5 100644
--- a/vm/src/builtins/type.rs
+++ b/vm/src/builtins/type.rs
@@ -8,13 +8,15 @@ use crate::common::{
     lock::{PyRwLock, PyRwLockReadGuard},
 };
 use crate::{
+    builtins::function::PyCellRef,
     builtins::PyBaseExceptionRef,
     class::{PyClassImpl, StaticType},
+    convert::ToPyObject,
     function::{FuncArgs, KwArgs, OptionalArg, PySetterValue},
     identifier,
     protocol::PyNumberMethods,
     types::{Callable, GetAttr, PyTypeFlags, PyTypeSlots, SetAttr},
-    AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
+    AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
 };
 use indexmap::{map::Entry, IndexMap};
 use itertools::Itertools;
@@ -622,6 +624,16 @@ impl PyType {
         let typ = Self::new_verbose_ref(name.as_str(), base, bases, attributes, slots, metatype)
             .map_err(|e| vm.new_type_error(e))?;
 
+        if let Some(cell) = typ.attributes.write().get(identifier!(vm, __classcell__)) {
+            let cell = PyCellRef::try_from_object(vm, cell.clone()).map_err(|_| {
+                vm.new_type_error(format!(
+                    "__classcell__ must be a nonlocal cell, not {}",
+                    cell.class().name()
+                ))
+            })?;
+            cell.set(Some(typ.clone().to_pyobject(vm)));
+        };
+
         // avoid deadlock
         let attributes = typ
             .attributes
diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs
index 8f3d99e47..70d5a142a 100644
--- a/vm/src/stdlib/builtins.rs
+++ b/vm/src/stdlib/builtins.rs
@@ -913,7 +913,19 @@ mod builtins {
         )?;
 
         if let Some(ref classcell) = classcell {
-            classcell.set(Some(class.clone()));
+            let classcell = classcell.get().ok_or_else(|| {
+                vm.new_type_error(format!(
+                    "__class__ not set defining {:?} as {:?}. Was __classcell__ propagated to type.__new__?",
+                    meta_name, class
+                ))
+            })?;
+
+            if !classcell.is(&class) {
+                return Err(vm.new_type_error(format!(
+                    "__class__ set to {:?} defining {:?} as {:?}",
+                    classcell, meta_name, class
+                )));
+            }
         }
 
         Ok(class)
diff --git a/vm/src/vm/context.rs b/vm/src/vm/context.rs
index 70df3070e..32f21f0f0 100644
--- a/vm/src/vm/context.rs
+++ b/vm/src/vm/context.rs
@@ -87,6 +87,7 @@ declare_const_name! {
     __ceil__,
     __cformat__,
     __class__,
+    __classcell__,
     __class_getitem__,
     __complex__,
     __contains__,