From 34a4abaf69ee065725d445dcf94ab52edd94f067 Mon Sep 17 00:00:00 2001 From: Dmitry Jemerov Date: Thu, 1 Jul 2010 17:30:04 +0400 Subject: [PATCH] recognize super calls via self.__class__ (PY-1190) --- python/src/com/jetbrains/python/PyNames.java | 1 + .../python/psi/impl/PyCallExpressionImpl.java | 61 ++++++++++++------- python/testData/resolve/SuperDunderClass.py | 10 +++ .../com/jetbrains/python/PyResolveTest.java | 4 ++ 4 files changed, 55 insertions(+), 21 deletions(-) create mode 100644 python/testData/resolve/SuperDunderClass.py diff --git a/python/src/com/jetbrains/python/PyNames.java b/python/src/com/jetbrains/python/PyNames.java index 23e16c899472..b350a3a57d25 100644 --- a/python/src/com/jetbrains/python/PyNames.java +++ b/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"; diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java index c2fcc3815fe1..76e7168dec52 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionImpl.java +++ b/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; + } } diff --git a/python/testData/resolve/SuperDunderClass.py b/python/testData/resolve/SuperDunderClass.py new file mode 100644 index 000000000000..daacee112388 --- /dev/null +++ b/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() +# + +B().foo() diff --git a/python/testSrc/com/jetbrains/python/PyResolveTest.java b/python/testSrc/com/jetbrains/python/PyResolveTest.java index 5908cb10ddf6..9db75e1fea8a 100644 --- a/python/testSrc/com/jetbrains/python/PyResolveTest.java +++ b/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"); + } } \ No newline at end of file