recognize super calls via self.__class__ (PY-1190)

python/src/com/jetbrains/python/PyNames.java

@@ -19,6 +19,7 @@ public class PyNames {
   @NonNls public static final String NEW = "__new__";
   @NonNls public static final String GETATTR = "__getattr__";
   @NonNls public static final String GETATTRIBUTE = "__getattribute__";
+  @NonNls public static final String CLASS = "__class__";
 
   @NonNls public static final String OBJECT = "object";
   @NonNls public static final String NONE = "None";

python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java

@@ -2,6 +2,7 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiReference;
 import com.intellij.psi.ResolveResult;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyNames;
@@ -129,6 +130,7 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress
     }
   }
 
+  @Nullable
   private PyType getSuperCallType(PyExpression callee, TypeEvalContext context) {
     PsiElement must_be_super_init = ((PyReferenceExpression)callee).getReference().resolve();
     if (must_be_super_init instanceof PyFunction) {
@@ -140,30 +142,23 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress
           if (args.length > 1) {
             PyExpression first_arg = args[0];
             if (first_arg instanceof PyReferenceExpression) {
-              PsiElement possible_class = ((PyReferenceExpression)first_arg).getReference().resolve();
-              if (possible_class instanceof PyClass && ((PyClass)possible_class).isNewStyleClass()) {
-                final PyClass first_class = (PyClass)possible_class;
-                // check 2nd argument, too; it should be an instance
-                PyExpression second_arg = args[1];
-                if (second_arg != null) {
-                  PyType second_type = second_arg.getType(context);
-                  if (second_type instanceof PyClassType) {
-                    // imitate isinstance(second_arg, possible_class)
-                    PyClass second_class = ((PyClassType)second_type).getPyClass();
-                    assert second_class != null;
-                    if (first_class == second_class) {
-                      final PyClass[] supers = first_class.getSuperClasses();
-                      if (supers.length > 0) {
-                        return new PyClassType(supers[0], false);
-                      }
-                    }
-                    if (second_class.isSubclass(first_class)) {
-                      // TODO: super(Foo, Bar) is a superclass of Foo directly preceding Bar in MRO
-                      return new PyClassType(first_class, false); // super(Foo, self) has type of Foo, modulo __get__()
-                    }
+              final PyReferenceExpression firstArgRef = (PyReferenceExpression)first_arg;
+              final PyExpression qualifier = firstArgRef.getQualifier();
+              if (qualifier != null && PyNames.CLASS.equals(firstArgRef.getReferencedName())) {
+                final PsiReference qRef = qualifier.getReference();
+                final PsiElement element = qRef == null ? null : qRef.resolve();
+                if (element instanceof PyParameter) {
+                  final PyParameterList parameterList = PsiTreeUtil.getParentOfType(element, PyParameterList.class);
+                  if (parameterList != null && element == parameterList.getParameters() [0]) {
+                    return getSuperCallType(context, PsiTreeUtil.getParentOfType(this, PyClass.class), args[1]);
                   }
                 }
               }
+              PsiElement possible_class = firstArgRef.getReference().resolve();
+              if (possible_class instanceof PyClass && ((PyClass)possible_class).isNewStyleClass()) {
+                final PyClass first_class = (PyClass)possible_class;
+                return getSuperCallType(context, first_class, args[1]);
+              }
             }
           }
         }
@@ -171,4 +166,28 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress
     }
     return null;
   }
+
+  @Nullable
+  private static PyType getSuperCallType(TypeEvalContext context, PyClass first_class, PyExpression second_arg) {
+    // check 2nd argument, too; it should be an instance
+    if (second_arg != null) {
+      PyType second_type = second_arg.getType(context);
+      if (second_type instanceof PyClassType) {
+        // imitate isinstance(second_arg, possible_class)
+        PyClass second_class = ((PyClassType)second_type).getPyClass();
+        assert second_class != null;
+        if (first_class == second_class) {
+          final PyClass[] supers = first_class.getSuperClasses();
+          if (supers.length > 0) {
+            return new PyClassType(supers[0], false);
+          }
+        }
+        if (second_class.isSubclass(first_class)) {
+          // TODO: super(Foo, Bar) is a superclass of Foo directly preceding Bar in MRO
+          return new PyClassType(first_class, false); // super(Foo, self) has type of Foo, modulo __get__()
+        }
+      }
+    }
+    return null;
+  }
 }

python/testData/resolve/SuperDunderClass.py

@@ -0,0 +1,10 @@
+class A(object):
+    def foo(self):
+        print "foo"
+
+class B(A):
+    def foo(self):
+        super(self.__class__, self).foo()
+#                                     <ref>
+
+B().foo()

python/testSrc/com/jetbrains/python/PyResolveTest.java

@@ -332,4 +332,8 @@ public class PyResolveTest extends PyResolveTestCase {
   public void testSuperMetaClass() {
     assertResolvesTo(PyFunction.class, "foo");
   }
+
+  public void testSuperDunderClass() {  // PY-1190
+    assertResolvesTo(PyFunction.class, "foo");
+  }
 }