diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java index 7ccb337eec7f..02da8b792d26 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java @@ -138,6 +138,7 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress if (must_be_super == PyBuiltinCache.getInstance(this).getClass(PyNames.SUPER)) { PyArgumentList arglist = getArgumentList(); if (arglist != null) { + final PyClass containingClass = PsiTreeUtil.getParentOfType(this, PyClass.class); PyExpression[] args = arglist.getArguments(); if (args.length > 1) { PyExpression first_arg = args[0]; @@ -150,7 +151,7 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress 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]); + return getSuperCallType(context, containingClass, args[1]); } } } @@ -161,6 +162,9 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress } } } + else if (((PyFile)getContainingFile()).getLanguageLevel().isPy3K()) { + return getFirstSuperClassType(containingClass); + } } } } @@ -177,10 +181,7 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress 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); - } + return getFirstSuperClassType(first_class); } if (second_class.isSubclass(first_class)) { // TODO: super(Foo, Bar) is a superclass of Foo directly preceding Bar in MRO @@ -190,4 +191,13 @@ public class PyCallExpressionImpl extends PyElementImpl implements PyCallExpress } return null; } + + private static PyType getFirstSuperClassType(PyClass first_class) { + // TODO handle __mro__ here + final PyClass[] supers = first_class.getSuperClasses(); + if (supers.length > 0) { + return new PyClassType(supers[0], false); + } + return null; + } } diff --git a/python/testData/resolve/SuperPy3k.py b/python/testData/resolve/SuperPy3k.py new file mode 100644 index 000000000000..dcb2185b1d56 --- /dev/null +++ b/python/testData/resolve/SuperPy3k.py @@ -0,0 +1,10 @@ +class A(object): + def foo(self): + print "foo" + +class B(A): + def foo(self): + super().foo() +# + +B().foo() diff --git a/python/testSrc/com/jetbrains/python/PyResolveTest.java b/python/testSrc/com/jetbrains/python/PyResolveTest.java index f4f4f2a8ce81..dbc2b8e74fd2 100644 --- a/python/testSrc/com/jetbrains/python/PyResolveTest.java +++ b/python/testSrc/com/jetbrains/python/PyResolveTest.java @@ -9,6 +9,7 @@ import com.intellij.psi.util.PsiTreeUtil; import com.jetbrains.python.fixtures.PyResolveTestCase; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyBuiltinCache; +import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher; import com.jetbrains.python.psi.resolve.ImportedResolveResult; public class PyResolveTest extends PyResolveTestCase { @@ -250,6 +251,17 @@ public class PyResolveTest extends PyResolveTestCase { assertEquals("A", ((PyFunction) targetElement).getContainingClass().getName()); } + public void testSuperPy3k() { // PY-1330 + PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), LanguageLevel.PYTHON30); + try { + final PyFunction pyFunction = assertResolvesTo(PyFunction.class, "foo"); + assertEquals("A", pyFunction.getContainingClass().getName()); + } + finally { + PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), null); + } + } + public void testStackOverflow() { PsiElement targetElement = resolve(); assertNull(targetElement); diff --git a/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java b/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java index bc597aff498a..4675856bf2b5 100644 --- a/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java +++ b/python/testSrc/com/jetbrains/python/fixtures/PyResolveTestCase.java @@ -47,11 +47,11 @@ public abstract class PyResolveTestCase extends PyLightFixtureTestCase { protected abstract PsiElement doResolve() throws Exception; - protected void assertResolvesTo(final Class aClass, final String name) { - assertResolvesTo(aClass, name, null); + protected T assertResolvesTo(final Class aClass, final String name) { + return assertResolvesTo(aClass, name, null); } - protected void assertResolvesTo(final Class aClass, + protected T assertResolvesTo(final Class aClass, final String name, String containingFilePath) { final PsiElement element; @@ -66,6 +66,7 @@ public abstract class PyResolveTestCase extends PyLightFixtureTestCase { if (containingFilePath != null) { assertEquals(containingFilePath, element.getContainingFile().getVirtualFile().getPath()); } + return (T)element; } protected int findMarkerOffset(final PsiFile psiFile) {