diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index b35be9ec5bc3..ad93e62f7f2c 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java @@ -251,26 +251,38 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { } @Nullable - private static PyFunctionType getSuperFunctionType(@NotNull PyFunction function, @NotNull TypeEvalContext context) { + private static PyFunction getOverriddenFunction(@NotNull PyFunction function, @NotNull TypeEvalContext context) { final Query superMethodSearchQuery = PySuperMethodsSearch.search(function, context); final PsiElement firstSuperMethod = superMethodSearchQuery.findFirst(); - if (firstSuperMethod instanceof PyFunction) { - final PyFunction superFunction = (PyFunction)firstSuperMethod; + if (!(firstSuperMethod instanceof PyFunction)) { + return null; + } + final PyFunction superFunction = (PyFunction)firstSuperMethod; - if (superFunction.getDecoratorList() != null) { - if (StreamEx.of(superFunction.getDecoratorList().getDecorators()) - .map(PyDecorator::getName) - .nonNull() - .anyMatch(PyNames.OVERLOAD::equals)) { - return null; - } + if (superFunction.getDecoratorList() != null) { + if (StreamEx.of(superFunction.getDecoratorList().getDecorators()) + .map(PyDecorator::getName) + .nonNull() + .anyMatch(PyNames.OVERLOAD::equals)) { + return null; } + } - final PyClass superClass = superFunction.getContainingClass(); - if (superClass != null && !PyNames.OBJECT.equals(superClass.getName())) { - return (PyFunctionType)context.getType(superFunction); - } + final PyClass superClass = superFunction.getContainingClass(); + if (superClass != null && !PyNames.OBJECT.equals(superClass.getName())) { + return superFunction; + } + + return null; + } + + + @Nullable + private static PyFunctionType getOverriddenFunctionType(@NotNull PyFunction function, @NotNull TypeEvalContext context) { + final PyFunction overriddenFunction = getOverriddenFunction(function, context); + if (overriddenFunction != null) { + return (PyFunctionType)context.getType(overriddenFunction); } return null; @@ -279,7 +291,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { @Nullable private static Ref getParameterTypeFromSupertype( @NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull TypeEvalContext context) { - final PyFunctionType superFunctionType = getSuperFunctionType(func, context); + final PyFunctionType superFunctionType = getOverriddenFunctionType(func, context); if (superFunctionType != null) { final PyFunctionType superFunctionTypeRemovedSelf = superFunctionType.dropSelf(context); @@ -350,12 +362,30 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { return null; } + /** + * Get function return type from supertype. + * + * The only source of type information in current implementation is annotation. This is to avoid false positives, + * that may arise from non direct type estimations (not from annotation, nor form type comments). + * + * TODO: switch to return type direct usage when type information source will be available. + * + * @param function + * @param context + * @return + */ @Nullable private static Ref getReturnTypeFromSupertype(@NotNull PyFunction function, @NotNull TypeEvalContext context) { - final PyFunctionType superFunctionType = getSuperFunctionType(function, context); + final PyFunction overriddenFunction = getOverriddenFunction(function, context); - if (superFunctionType != null) { - return new Ref<>(context.getReturnType(superFunctionType.getCallable())); + if (overriddenFunction != null) { + PyExpression superFunctionAnnotation = getReturnTypeAnnotation(overriddenFunction, context); + if (superFunctionAnnotation != null) { + final Ref typeRef = getType(superFunctionAnnotation, new Context(context)); + if (typeRef != null) { + return Ref.create(toAsyncIfNeeded(function, typeRef.get())); + } + } } return null; diff --git a/python/testSrc/com/jetbrains/python/Py3TypeTest.java b/python/testSrc/com/jetbrains/python/Py3TypeTest.java index 81cc4ff9eef9..de20eca28a17 100644 --- a/python/testSrc/com/jetbrains/python/Py3TypeTest.java +++ b/python/testSrc/com/jetbrains/python/Py3TypeTest.java @@ -853,7 +853,12 @@ public class Py3TypeTest extends PyTestCase { ); } - public void testReturnTypeInferenceInSubclassFromDocstring() { + /** + * TODO: activate when return type information from :rtype: will be available in subclasses. + * + * See {@code {@link com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider#getReturnTypeFromSupertype}} javadoc. + */ + public void ignoreTestReturnTypeInferenceInSubclassFromDocstring() { runWithLanguageLevel( LanguageLevel.PYTHON35, () -> doTest("int", @@ -871,7 +876,12 @@ public class Py3TypeTest extends PyTestCase { ); } - public void testReturnTypeInferenceInSubclassHierarchyFromDocstring() { + /** + * TODO: activate when return type information from :rtype: will be available in subclasses. + * + * See {@code {@link com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider#getReturnTypeFromSupertype}} javadoc. + */ + public void ignoreTestReturnTypeInferenceInSubclassHierarchyFromDocstring() { runWithLanguageLevel( LanguageLevel.PYTHON35, () -> doTest("int",