Merge pull request #698 from palaviv/super

windelbouwman · web-flow · commit 7ed4bc4aa8c8 · 2019-03-22T17:13:13.000+01:00
Use first argument in super
diff --git a/tests/snippets/class.py b/tests/snippets/class.py
@@ -75,6 +75,33 @@ def test1(self):
 assert c.test() == 100
 assert c.test1() == 200
 
+class Me():
+
+    def test(me):
+        return 100
+
+class Me2(Me):
+
+    def test(me):
+        return super().test()
+
+class A():
+    def f(self):
+        pass
+
+class B(A):
+    def f(self):
+        super().f()
+
+class C(B):
+    def f(self):
+        super().f()
+
+C().f()
+
+me = Me2()
+assert me.test() == 100
+
 a = super(bool, True)
 assert isinstance(a, super)
 assert type(a) is super
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -126,6 +126,7 @@ pub trait NameProtocol {
     fn load_name(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef>;
     fn store_name(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef);
     fn delete_name(&self, vm: &VirtualMachine, name: &str);
+    fn load_cell(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef>;
 }
 
 impl NameProtocol for Scope {
@@ -143,6 +144,15 @@ impl NameProtocol for Scope {
         vm.builtins.get_item(name)
     }
 
+    fn load_cell(&self, _vm: &VirtualMachine, name: &str) -> Option<PyObjectRef> {
+        for dict in self.locals.iter().skip(1) {
+            if let Some(value) = dict.get_item(name) {
+                return Some(value);
+            }
+        }
+        None
+    }
+
     fn store_name(&self, vm: &VirtualMachine, key: &str, value: PyObjectRef) {
         self.get_locals().set_item(&vm.ctx, key, value)
     }
diff --git a/vm/src/obj/objsuper.rs b/vm/src/obj/objsuper.rs
@@ -6,7 +6,9 @@ https://github.com/python/cpython/blob/50b48572d9a90c5bb36e2bef6179548ea927a35a/
 
 */
 
+use crate::frame::NameProtocol;
 use crate::function::PyFuncArgs;
+use crate::obj::objstr;
 use crate::obj::objtype::PyClass;
 use crate::pyobject::{
     DictProtocol, PyContext, PyObject, PyObjectRef, PyResult, PyValue, TypeProtocol,
@@ -18,6 +20,7 @@ use super::objtype;
 #[derive(Debug)]
 pub struct PySuper {
     obj: PyObjectRef,
+    typ: PyObjectRef,
 }
 
 impl PyValue for PySuper {
@@ -67,15 +70,20 @@ fn super_getattribute(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     );
 
     let inst = super_obj.payload::<PySuper>().unwrap().obj.clone();
+    let typ = super_obj.payload::<PySuper>().unwrap().typ.clone();
 
-    match inst.typ().payload::<PyClass>() {
+    match 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)))
+            Err(vm.new_attribute_error(format!(
+                "{} has no attribute '{}'",
+                inst,
+                objstr::get_value(name_str)
+            )))
         }
         _ => panic!("not Class"),
     }
@@ -98,9 +106,13 @@ fn super_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     let py_type = if let Some(ty) = py_type {
         ty.clone()
     } else {
-        match vm.get_locals().get_item("self") {
-            Some(obj) => obj.typ().clone(),
-            _ => panic!("No self"),
+        match vm.current_scope().load_cell(vm, "__class__") {
+            Some(obj) => obj.clone(),
+            _ => {
+                return Err(vm.new_type_error(
+                    "super must be called with 1 argument or from inside class method".to_string(),
+                ));
+            }
         }
     };
 
@@ -117,9 +129,19 @@ fn super_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     let py_obj = if let Some(obj) = py_obj {
         obj.clone()
     } else {
-        match vm.get_locals().get_item("self") {
-            Some(obj) => obj,
-            _ => panic!("No self"),
+        let frame = vm.current_frame();
+        if let Some(first_arg) = frame.code.arg_names.get(0) {
+            match vm.get_locals().get_item(first_arg) {
+                Some(obj) => obj.clone(),
+                _ => {
+                    return Err(vm
+                        .new_type_error(format!("super arguement {} was not supplied", first_arg)));
+                }
+            }
+        } else {
+            return Err(vm.new_type_error(
+                "super must be called with 1 argument or from inside class method".to_string(),
+            ));
         }
     };
 
@@ -130,5 +152,11 @@ fn super_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
         ));
     }
 
-    Ok(PyObject::new(PySuper { obj: py_obj }, cls.clone()))
+    Ok(PyObject::new(
+        PySuper {
+            obj: py_obj,
+            typ: py_type,
+        },
+        cls.clone(),
+    ))
 }
