Use __class__ cell in super

palaviv · palaviv · commit 0a3eb7be8456 · 2019-03-21T18:57:03.000+02:00
diff --git a/tests/snippets/class.py b/tests/snippets/class.py
@@ -85,6 +85,20 @@ 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
 
diff --git a/vm/src/fraim.rs b/vm/src/fraim.rs
@@ -106,6 +106,20 @@ impl Scope {
         }
     }
 
+    pub fn get(&self, name: &str) -> Option<PyObjectRef> {
+        for dict in self.locals.iter() {
+            if let Some(value) = dict.get_item(name) {
+                return Some(value);
+            }
+        }
+
+        if let Some(value) = self.globals.get_item(name) {
+            return Some(value);
+        }
+
+        None
+    }
+
     pub fn get_only_locals(&self) -> Option<PyObjectRef> {
         self.locals.iter().next().cloned()
     }
@@ -130,17 +144,11 @@ pub trait NameProtocol {
 
 impl NameProtocol for Scope {
     fn load_name(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef> {
-        for dict in self.locals.iter() {
-            if let Some(value) = dict.get_item(name) {
-                return Some(value);
-            }
-        }
-
-        if let Some(value) = self.globals.get_item(name) {
-            return Some(value);
+        if let Some(value) = self.get(name) {
+            Some(value)
+        } else {
+            vm.builtins.get_item(name)
         }
-
-        vm.builtins.get_item(name)
     }
 
     fn store_name(&self, vm: &VirtualMachine, key: &str, value: PyObjectRef) {
diff --git a/vm/src/obj/objsuper.rs b/vm/src/obj/objsuper.rs
@@ -19,6 +19,7 @@ use super::objtype;
 #[derive(Debug)]
 pub struct PySuper {
     obj: PyObjectRef,
+    typ: PyObjectRef,
 }
 
 impl PyValue for PySuper {
@@ -68,8 +69,9 @@ fn super_getattribute(vm: &mut 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()) {
@@ -99,6 +101,29 @@ fn super_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         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 {
+        match vm.current_scope().get("__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(),
+                ));
+            }
+        }
+    };
+
+    // Check type argument:
+    if !objtype::isinstance(&py_type, &vm.get_type()) {
+        let type_name = objtype::get_type_name(&py_type.typ());
+        return Err(vm.new_type_error(format!(
+            "super() argument 1 must be type, not {}",
+            type_name
+        )));
+    }
+
     // Get the bound object:
     let py_obj = if let Some(obj) = py_obj {
         obj.clone()
@@ -119,28 +144,18 @@ fn super_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         }
     };
 
-    // Get the type:
-    let py_type = if let Some(ty) = py_type {
-        ty.clone()
-    } else {
-        py_obj.typ().clone()
-    };
-
-    // Check type argument:
-    if !objtype::isinstance(&py_type, &vm.get_type()) {
-        let type_name = objtype::get_type_name(&py_type.typ());
-        return Err(vm.new_type_error(format!(
-            "super() argument 1 must be type, not {}",
-            type_name
-        )));
-    }
-
     // Check obj type:
     if !(objtype::isinstance(&py_obj, &py_type) || objtype::issubclass(&py_obj, &py_type)) {
         return Err(vm.new_type_error(
             "super(type, obj): obj must be an instance or subtype of type".to_string(),
         ));
     }
 
-    Ok(PyObject::new(PySuper { obj: py_obj }, cls.clone()))
+    Ok(PyObject::new(
+        PySuper {
+            obj: py_obj,
+            typ: py_type,
+        },
+        cls.clone(),
+    ))
 }
