Merge pull request RustPython#185 from RustPython/prepare

windelbouwman · web-flow · commit 081a33f41aac · 2018-11-03T10:25:59.000+01:00
Add support for __prepare__.
diff --git a/tests/snippets/metaclasses.py b/tests/snippets/metaclasses.py
@@ -2,6 +2,9 @@ class MC(type):
     classes = []
     count = 0
 
+    def __prepare__(name, bases):
+        return {'prepared': True}
+
     def __new__(cls, name, bases, namespace):
         MC.classes.append(name)
         return type.__new__(cls, name, bases, namespace)
@@ -16,10 +19,13 @@ def __new__(cls, count):
         self.count = count
         return self
 
-class D(object, metaclass=MC):
+class D(metaclass=MC):
     pass
 
 assert MC == type(C)
 assert C == type(C())
 assert MC.classes == ['C', 'D']
 assert C().count == 2
+
+assert C.prepared
+assert D.prepared
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -635,7 +635,19 @@ pub fn builtin_build_class_(vm: &mut VirtualMachine, mut args: PyFuncArgs) -> Py
     let metaclass = args.get_kwarg("metaclass", vm.get_type());
 
     bases.push(vm.context().object());
-    let namespace = vm.new_dict();
+    let bases = vm.context().new_tuple(bases);
+
+    // Prepare uses full __getattribute__ resolution chain.
+    let prepare_name = vm.new_str("__prepare__".to_string());
+    let prepare = vm.get_attribute(metaclass.clone(), prepare_name)?;
+    let namespace = vm.invoke(
+        prepare,
+        PyFuncArgs {
+            args: vec![name_arg.clone(), bases.clone()],
+            kwargs: vec![],
+        },
+    )?;
+
     &vm.invoke(
         function,
         PyFuncArgs {
@@ -644,8 +656,6 @@ pub fn builtin_build_class_(vm: &mut VirtualMachine, mut args: PyFuncArgs) -> Py
         },
     );
 
-    let bases = vm.context().new_tuple(bases);
-
     // Special case: __new__ must be looked up on the metaclass, not the meta-metaclass as
     // per vm.call(metaclass, "__new__", ...)
     let new = metaclass.get_attr("__new__").unwrap();
diff --git a/vm/src/obj/objtype.rs b/vm/src/obj/objtype.rs
@@ -28,6 +28,7 @@ pub fn init(context: &PyContext) {
     type_type.set_attr("__mro__", context.new_member_descriptor(type_mro));
     type_type.set_attr("__class__", context.new_member_descriptor(type_new));
     type_type.set_attr("__repr__", context.new_rustfunc(type_repr));
+    type_type.set_attr("__prepare__", context.new_rustfunc(type_prepare));
 }
 
 fn type_mro(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -291,6 +292,10 @@ fn type_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.new_str(format!("<class '{}'>", type_name)))
 }
 
+fn type_prepare(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult {
+    Ok(vm.new_dict())
+}
+
 #[cfg(test)]
 mod tests {
     use super::{linearise_mro, new};
