Merge pull request #696 from palaviv/super

super in class methods

tests/snippets/class.py

@@ -50,12 +50,30 @@ class Bar2(Bar):
     def __init__(self):
         super().__init__(101)
 
+bar2 = Bar2()
+assert bar2.get_x() == 101
 
-# TODO: make this work:
-# bar2 = Bar2()
-# assert bar2.get_x() == 101
+class A():
+    def test(self):
+        return 100
 
-a = super(int, 2)
+class B():
+    def test1(self):
+        return 200
+
+class C(A,B):
+    def test(self):
+        return super().test()
+
+    def test1(self):
+        return super().test1()
+
+c = C()
+assert c.test() == 100
+assert c.test1() == 200
+
+a = super(bool, True)
 assert isinstance(a, super)
 assert type(a) is super
+assert a.conjugate() == 1
 

vm/src/obj/objsuper.rs

@@ -7,11 +7,25 @@ https://github.com/python/cpython/blob/50b48572d9a90c5bb36e2bef6179548ea927a35a/
 */
 
 use crate::function::PyFuncArgs;
-use crate::pyobject::{PyContext, PyResult, TypeProtocol};
+use crate::obj::objtype::PyClass;
+use crate::pyobject::{
+    DictProtocol, PyContext, PyObject, PyObjectRef, PyResult, PyValue, TypeProtocol,
+};
 use crate::vm::VirtualMachine;
 
 use super::objtype;
 
+#[derive(Debug)]
+pub struct PySuper {
+    obj: PyObjectRef,
+}
+
+impl PyValue for PySuper {
+    fn class(vm: &mut VirtualMachine) -> PyObjectRef {
+        vm.ctx.super_type()
+    }
+}
+
 pub fn init(context: &PyContext) {
     let super_type = &context.super_type;
 
@@ -29,7 +43,12 @@ pub fn init(context: &PyContext) {
                      def cmeth(cls, arg):\n        \
                      super().cmeth(arg)\n";
 
-    context.set_attr(&super_type, "__init__", context.new_rustfunc(super_init));
+    context.set_attr(&super_type, "__new__", context.new_rustfunc(super_new));
+    context.set_attr(
+        &super_type,
+        "__getattribute__",
+        context.new_rustfunc(super_getattribute),
+    );
     context.set_attr(
         &super_type,
         "__doc__",
@@ -37,24 +56,52 @@ pub fn init(context: &PyContext) {
     );
 }
 
-fn super_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    trace!("super.__init__ {:?}", args.args);
+fn super_getattribute(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
         args,
-        required = [(inst, None)],
+        required = [
+            (super_obj, Some(vm.ctx.super_type())),
+            (name_str, Some(vm.ctx.str_type()))
+        ]
+    );
+
+    let inst = super_obj.payload::<PySuper>().unwrap().obj.clone();
+
+    match inst.typ().payload::<PyClass>() {
+        Some(PyClass { ref mro, .. }) => {
+            for class in mro {
+                if let Ok(item) = vm.get_attribute(class.as_object().clone(), name_str.clone()) {
+                    return Ok(vm.ctx.new_bound_method(item, inst.clone()));
+                }
+            }
+            Err(vm.new_attribute_error(format!("{} has no attribute '{}'", inst, name_str)))
+        }
+        _ => panic!("not Class"),
+    }
+}
+
+fn super_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    trace!("super.__new__ {:?}", args.args);
+    arg_check!(
+        vm,
+        args,
+        required = [(cls, None)],
         optional = [(py_type, None), (py_obj, None)]
     );
 
+    if !objtype::issubclass(cls, &vm.ctx.super_type()) {
+        return Err(vm.new_type_error(format!("{:?} is not a subtype of super", cls)));
+    }
+
     // Get the type:
     let py_type = if let Some(ty) = py_type {
         ty.clone()
     } else {
-        // TODO: implement complex logic here....
-        unimplemented!("TODO: get frame and determine instance and class?");
-        // let frame = vm.get_current_frame();
-        //
-        // vm.get_none()
+        match vm.get_locals().get_item("self") {
+            Some(obj) => obj.typ().clone(),
+            _ => panic!("No self"),
+        }
     };
 
     // Check type argument:
@@ -70,7 +117,10 @@ fn super_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let py_obj = if let Some(obj) = py_obj {
         obj.clone()
     } else {
-        vm.get_none()
+        match vm.get_locals().get_item("self") {
+            Some(obj) => obj,
+            _ => panic!("No self"),
+        }
     };
 
     // Check obj type:
@@ -80,9 +130,5 @@ fn super_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         ));
     }
 
-    // TODO: how to store those types?
-    vm.ctx.set_attr(inst, "type", py_type);
-    vm.ctx.set_attr(inst, "obj", py_obj);
-
-    Ok(vm.get_none())
+    Ok(PyObject::new(PySuper { obj: py_obj }, cls.clone()))
 }