__class_getitem__ isn't a special method

Author: coolreader18

vm/src/builtins/make_module.rs

@@ -20,7 +20,7 @@ mod decl {
     use crate::builtins::pybool::IntoPyBool;
     use crate::builtins::pystr::{PyStr, PyStrRef};
     use crate::builtins::pytype::PyTypeRef;
-    use crate::builtins::{PyByteArray, PyBytes};
+    use crate::builtins::{PyByteArray, PyBytes, PyTupleRef};
     use crate::byteslike::ArgBytesLike;
     use crate::common::{hash::PyHash, str::to_ascii};
     #[cfg(feature = "rustpython-compiler")]
@@ -832,7 +832,7 @@ mod decl {
     pub fn __build_class__(
         function: PyFunctionRef,
         qualified_name: PyStrRef,
-        bases: Args<PyTypeRef>,
+        bases: Args,
         mut kwargs: KwArgs,
         vm: &VirtualMachine,
     ) -> PyResult {
@@ -845,7 +845,41 @@ mod decl {
             vm.ctx.types.type_type.clone()
         };
 
-        for base in bases.clone() {
+        let mut new_bases: Option<Vec<PyObjectRef>> = None;
+
+        let bases = PyTupleRef::with_elements(bases.into_vec(), &vm.ctx);
+
+        for (i, base) in bases.as_slice().iter().enumerate() {
+            if base.isinstance(&vm.ctx.types.type_type) {
+                if let Some(bases) = &mut new_bases {
+                    bases.push(base.clone());
+                }
+                continue;
+            }
+            let mro_entries = vm.get_attribute_opt(base.clone(), "__mro_entries__")?;
+            let entries = match mro_entries {
+                Some(meth) => vm.invoke(&meth, (bases.clone(),))?,
+                None => {
+                    if let Some(bases) = &mut new_bases {
+                        bases.push(base.clone());
+                    }
+                    continue;
+                }
+            };
+            let entries: PyTupleRef = entries
+                .downcast()
+                .map_err(|_| vm.new_type_error("__mro_entries__ must return a tuple".to_owned()))?;
+            let new_bases = new_bases.get_or_insert_with(|| bases.as_slice()[..i].to_vec());
+            new_bases.extend_from_slice(entries.as_slice());
+        }
+
+        let new_bases = new_bases.map(|v| PyTupleRef::with_elements(v, &vm.ctx));
+        let (orig_bases, bases) = match new_bases {
+            Some(new) => (Some(bases), new),
+            None => (None, bases),
+        };
+
+        for base in bases.as_slice().iter() {
             let base_class = base.class();
             if base_class.issubclass(&metaclass) {
                 metaclass = base.clone_class();
@@ -858,7 +892,7 @@ mod decl {
             }
         }
 
-        let bases = bases.into_tuple(vm);
+        let bases = bases.into_object();
 
         // Prepare uses full __getattribute__ resolution chain.
         let prepare = vm.get_attribute(metaclass.clone().into_object(), "__prepare__")?;
@@ -872,6 +906,10 @@ mod decl {
         let classcell = function.invoke_with_locals(().into(), Some(namespace.clone()), vm)?;
         let classcell = <Option<PyCellRef>>::try_from_object(vm, classcell)?;
 
+        if let Some(orig_bases) = orig_bases {
+            namespace.set_item("__orig_bases__", orig_bases.into_object(), vm)?;
+        }
+
         let class = vm.invoke(
             metaclass.as_object(),
             FuncArgs::new(vec![name_obj, bases, namespace.into_object()], kwargs),

vm/src/pyobject.rs

@@ -686,25 +686,19 @@ where
 {
     fn get_item(&self, key: T, vm: &VirtualMachine) -> PyResult {
         match vm.get_special_method(self.clone(), "__getitem__")? {
-            Ok(special_method) => special_method.invoke((key,), vm),
+            Ok(special_method) => return special_method.invoke((key,), vm),
             Err(obj) => {
                 if obj.isinstance(&vm.ctx.types.type_type) {
-                    vm.get_special_method(obj, "__class_getitem__")?
-                        .map_err(|obj2| {
-                            vm.new_type_error(format!(
-                                "'{}' object is not subscriptable",
-                                obj2.class().name
-                            ))
-                        })?
-                        .invoke((key,), vm)
-                } else {
-                    Err(vm.new_type_error(format!(
-                        "'{}' object is not subscriptable",
-                        obj.class().name
-                    )))
+                    if let Some(class_getitem) = vm.get_attribute_opt(obj, "__class_getitem__")? {
+                        return vm.invoke(&class_getitem, (key,));
+                    }
                 }
             }
         }
+        Err(vm.new_type_error(format!(
+            "'{}' object is not subscriptable",
+            self.class().name
+        )))
     }
 
     fn set_item(&self, key: T, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {