Rework tp_new

vm/src/macros.rs

@@ -153,6 +153,7 @@ macro_rules! extend_class {
         $(
             $crate::extend_class!(@set_attr($ctx, $class, $name, $value));
         )*
+        $ctx.add_tp_new_wrapper(&$class);
     };
 
     (@set_attr($ctx:expr, $class:expr, (slot $slot_name:ident), $value:expr)) => {

vm/src/obj/objobject.rs

@@ -2,7 +2,7 @@ use super::objbool;
 use super::objdict::PyDictRef;
 use super::objlist::PyList;
 use super::objstr::PyStringRef;
-use super::objtype::{self, PyClassRef};
+use super::objtype::PyClassRef;
 use crate::function::{OptionalArg, PyFuncArgs};
 use crate::pyhash;
 use crate::pyobject::{
@@ -262,10 +262,6 @@ pub(crate) fn setattr(
 
 pub fn init(context: &PyContext) {
     PyBaseObject::extend_class(context, &context.types.object_type);
-    extend_class!(context, &context.types.object_type, {
-        // yeah, it's `type_new`, but we're putting here so it's available on every object
-        "__new__" => context.new_classmethod(objtype::type_new),
-    });
 }
 
 fn common_reduce(obj: PyObjectRef, proto: usize, vm: &VirtualMachine) -> PyResult {

vm/src/obj/objtype.rs

@@ -5,6 +5,7 @@ use std::fmt;
 use super::objdict::PyDictRef;
 use super::objlist::PyList;
 use super::objmappingproxy::PyMappingProxy;
+use super::objstaticmethod::PyStaticMethod;
 use super::objstr::PyStringRef;
 use super::objtuple::PyTuple;
 use super::objweakref::PyWeak;
@@ -48,7 +49,7 @@ impl PyValue for PyClass {
 impl PyClassRef {
     pub fn iter_mro<'a>(&'a self) -> impl Iterator<Item = &'a PyClassRef> + DoubleEndedIterator {
         std::iter::once(self).chain(self.mro.iter())
-        }
+    }
 
     #[pyproperty(name = "__mro__")]
     fn get_mro(self) -> PyTuple {
@@ -282,9 +283,16 @@ impl PyClassRef {
             (metatype, base, bases)
         };
 
-        let attributes = dict.to_attributes();
+        let mut attributes = dict.to_attributes();
+        if let Some(f) = attributes.get_mut("__new__") {
+            if f.class().is(&vm.ctx.function_type()) {
+                *f = PyStaticMethod::new(f.clone()).into_ref(vm).into_object();
+            }
+        }
+
         let typ = new(metatype, name.as_str(), base.clone(), bases, attributes)?;
         typ.slots.borrow_mut().flags = base.slots.borrow().flags;
+        vm.ctx.add_tp_new_wrapper(&typ);
         Ok(typ.into())
     }
 
@@ -292,9 +300,12 @@ impl PyClassRef {
     #[pymethod(magic)]
     fn call(self, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult {
         vm_trace!("type_call: {:?}", self);
-        let new = vm.get_attribute(self.as_object().clone(), "__new__")?;
-        let new_args = args.insert(self.into_object());
-        let obj = vm.invoke(&new, new_args)?;
+        let obj = call_tp_new(self.clone(), self.clone(), args.clone(), vm)?;
+
+        if (self.is(&vm.ctx.types.type_type) && args.kwargs.is_empty()) || !isinstance(&obj, &self)
+        {
+            return Ok(obj);
+        }
 
         if let Some(init_method_or_err) = vm.get_method(obj.clone(), "__init__") {
             let init_method = init_method_or_err?;
@@ -340,9 +351,32 @@ pub fn isinstance<T: TypeProtocol>(obj: &T, cls: &PyClassRef) -> bool {
 pub fn issubclass(subclass: &PyClassRef, cls: &PyClassRef) -> bool {
     subclass.iter_mro().any(|c| c.is(cls))
 }
+
+fn call_tp_new(
+    typ: PyClassRef,
+    subtype: PyClassRef,
+    args: PyFuncArgs,
+    vm: &VirtualMachine,
+) -> PyResult {
+    for cls in typ.iter_mro() {
+        if let Some(new_meth) = cls.get_attr("__new__") {
+            if !vm.ctx.is_tp_new_wrapper(&new_meth) {
+                let new_meth = vm.call_if_get_descriptor(new_meth, typ.clone().into_object())?;
+                return vm.invoke(&new_meth, args.insert(typ.clone().into_object()));
+            }
+        }
+    }
+    let class_with_new_slot = typ
+        .iter_mro()
+        .cloned()
+        .find(|cls| cls.slots.borrow().new.is_some())
+        .expect("Should be able to find a new slot somewhere in the mro");
+    let slots = class_with_new_slot.slots.borrow();
+    let new_slot = slots.new.as_ref().unwrap();
+    new_slot(vm, args.insert(subtype.into_object()))
 }
 
-pub fn type_new(
+pub fn tp_new_wrapper(
     zelf: PyClassRef,
     cls: PyClassRef,
     args: PyFuncArgs,
@@ -355,21 +389,7 @@ pub fn type_new(
             cls = cls.name,
         )));
     }
-
-    let class_with_new_slot = if cls.slots.borrow().new.is_some() {
-        cls.clone()
-    } else {
-        cls.mro
-            .iter()
-            .cloned()
-            .find(|cls| cls.slots.borrow().new.is_some())
-            .expect("Should be able to find a new slot somewhere in the mro")
-    };
-
-    let slots = class_with_new_slot.slots.borrow();
-    let new = slots.new.as_ref().unwrap();
-
-    new(vm, args.insert(cls.into_object()))
+    call_tp_new(zelf, cls, args, vm)
 }
 
 impl PyClassRef {
@@ -393,7 +413,7 @@ impl PyClassRef {
     // This is the internal has_attr implementation for fast lookup on a class.
     pub fn has_attr(&self, attr_name: &str) -> bool {
         self.iter_mro()
-                .any(|c| c.attributes.borrow().contains_key(attr_name))
+            .any(|c| c.attributes.borrow().contains_key(attr_name))
     }
 
     pub fn get_attributes(self) -> PyAttributes {

vm/src/pyobject.rs

@@ -106,6 +106,7 @@ pub struct PyContext {
     pub types: TypeZoo,
     pub exceptions: exceptions::ExceptionZoo,
     pub int_cache_pool: Vec<PyObjectRef>,
+    tp_new_wrapper: PyObjectRef,
 }
 
 pub type PyNotImplementedRef = PyRef<PyNotImplemented>;
@@ -160,6 +161,12 @@ impl PyContext {
 
         let empty_tuple = create_object(PyTuple::from(vec![]), &types.tuple_type);
 
+        let tp_new_wrapper = create_object(
+            PyBuiltinFunction::new(objtype::tp_new_wrapper.into_func()),
+            &types.builtin_function_or_method_type,
+        )
+        .into_object();
+
         let context = PyContext {
             true_value,
             false_value,
@@ -172,6 +179,7 @@ impl PyContext {
             types,
             exceptions,
             int_cache_pool,
+            tp_new_wrapper,
         };
         initialize_types(&context);
 
@@ -592,6 +600,19 @@ impl PyContext {
             bytecode::Constant::Ellipsis => self.ellipsis(),
         }
     }
+
+    pub fn add_tp_new_wrapper(&self, ty: &PyClassRef) {
+        if !ty.attributes.borrow().contains_key("__new__") {
+            let new_wrapper =
+                self.new_bound_method(self.tp_new_wrapper.clone(), ty.clone().into_object());
+            ty.set_str_attr("__new__", new_wrapper);
+        }
+    }
+
+    pub fn is_tp_new_wrapper(&self, obj: &PyObjectRef) -> bool {
+        obj.payload::<PyBoundMethod>()
+            .map_or(false, |bound| bound.function.is(&self.tp_new_wrapper))
+    }
 }
 
 impl Default for PyContext {
@@ -1231,6 +1252,7 @@ pub trait PyClassImpl: PyClassDef {
     fn extend_class(ctx: &PyContext, class: &PyClassRef) {
         Self::impl_extend_class(ctx, class);
         class.slots.borrow_mut().flags = Self::TP_FLAGS;
+        ctx.add_tp_new_wrapper(&class);
         if let Some(doc) = Self::DOC {
             class.set_str_attr("__doc__", ctx.new_str(doc.into()));
         }