From 4b7f436e13ebf322fd303aa6bfa7ef70285ab95b Mon Sep 17 00:00:00 2001 From: Semyon Proshev Date: Fri, 12 May 2017 19:19:00 +0300 Subject: [PATCH] PY-22971 Fixed: Support @typing.overload in regular Python files, not only in Python stubs Don't take implementation into account while inferring call type if there is at least one overload. --- .../python/PyParameterInfoHandler.java | 4 +- .../inspections/PyArgumentListInspection.java | 2 +- .../psi/impl/PyCallExpressionHelper.java | 62 ++++--- .../python/psi/types/PyTypeChecker.java | 3 +- .../jetbrains/python/pyi/PyiTypeProvider.java | 5 +- .../b/__init__.py | 14 ++ .../b/__init__.py | 15 ++ .../b/__init__.py | 14 ++ .../b/__init__.py | 15 ++ .../b/__init__.py | 14 ++ .../b/__init__.py | 15 ++ .../com/jetbrains/python/PyTypeTest.java | 171 ++++++++++++++++++ 12 files changed, 300 insertions(+), 34 deletions(-) create mode 100644 python/testData/types/FirstOverloadAndImplementationInImportedClass/b/__init__.py create mode 100644 python/testData/types/FirstOverloadAndImplementationInImportedModule/b/__init__.py create mode 100644 python/testData/types/NotMatchedOverloadsAndImplementationInImportedClass/b/__init__.py create mode 100644 python/testData/types/NotMatchedOverloadsAndImplementationInImportedModule/b/__init__.py create mode 100644 python/testData/types/SecondOverloadAndImplementationInImportedClass/b/__init__.py create mode 100644 python/testData/types/SecondOverloadAndImplementationInImportedModule/b/__init__.py diff --git a/python/src/com/jetbrains/python/PyParameterInfoHandler.java b/python/src/com/jetbrains/python/PyParameterInfoHandler.java index 28cca3fe0bbb..5c4076b4b536 100644 --- a/python/src/com/jetbrains/python/PyParameterInfoHandler.java +++ b/python/src/com/jetbrains/python/PyParameterInfoHandler.java @@ -22,6 +22,7 @@ import com.intellij.openapi.util.Pair; import com.intellij.openapi.util.TextRange; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; +import com.intellij.psi.ResolveResult; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.ArrayUtil; import com.intellij.util.text.CharArrayUtil; @@ -75,7 +76,8 @@ public class PyParameterInfoHandler implements ParameterInfoHandler ratedMarkedCallees = PyUtil.filterTopPriorityResults(call.multiResolveRatedCallee(resolveContext)); - final Object[] items = PyCallExpressionHelper.forEveryScopeTakeOverloadsOtherwiseImplementations(ratedMarkedCallees, typeEvalContext) + final Object[] items = PyCallExpressionHelper + .forEveryScopeTakeOverloadsOtherwiseImplementations(ratedMarkedCallees, ResolveResult::getElement, typeEvalContext) .map(ratedMarkedCallee -> Pair.createNonNull(call, ratedMarkedCallee.getMarkedCallee())) .toArray(); diff --git a/python/src/com/jetbrains/python/inspections/PyArgumentListInspection.java b/python/src/com/jetbrains/python/inspections/PyArgumentListInspection.java index 345bcef9938f..b148c78eb811 100644 --- a/python/src/com/jetbrains/python/inspections/PyArgumentListInspection.java +++ b/python/src/com/jetbrains/python/inspections/PyArgumentListInspection.java @@ -155,7 +155,7 @@ public class PyArgumentListInspection extends PyInspection { PyUtil.filterTopPriorityResults(call.multiResolveRatedCallee(resolveContext, implicitOffset)); return PyCallExpressionHelper - .forEveryScopeTakeOverloadsOtherwiseImplementations(ratedMarkedCallees, context) + .forEveryScopeTakeOverloadsOtherwiseImplementations(ratedMarkedCallees, ResolveResult::getElement, context) .map(ratedMarkedCallee -> PyCallExpressionHelper.mapArguments(call, ratedMarkedCallee.getMarkedCallee(), context)) .collect(Collectors.toList()); } diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java index 645cdc5cc05b..9fe41fb6a917 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java +++ b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java @@ -19,8 +19,8 @@ import com.intellij.codeInsight.completion.CompletionUtil; import com.intellij.openapi.util.Pair; import com.intellij.openapi.util.Ref; import com.intellij.psi.PsiElement; -import com.intellij.psi.PsiPolyVariantReference; import com.intellij.psi.PsiReference; +import com.intellij.psi.ResolveResult; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.ObjectUtils; import com.intellij.util.containers.ContainerUtil; @@ -39,6 +39,7 @@ import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import java.util.*; +import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -394,17 +395,19 @@ public class PyCallExpressionHelper { } // normal cases final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); - final PsiPolyVariantReference reference = ((PyReferenceExpression)callee).getReference(resolveContext); - final List members = new ArrayList<>(); - for (PsiElement target : PyUtil.multiResolveTopPriority(reference)) { - PyUtil.verboseOnly(() ->PyPsiUtils.assertValid(target)); - if (target != null) { - final Ref typeRef = getCallTargetReturnType(call, target, context); - if (typeRef != null) { - members.add(typeRef.get()); - } - } - } + final List resolveResults = + PyUtil.filterTopPriorityResults(multiResolveCallee(callee, resolveContext)); + + final List members = StreamEx + .of(forEveryScopeTakeOverloadsOtherwiseImplementations(resolveResults, ResolveResult::getElement, context)) + .map(ResolveResult::getElement) + .nonNull() + .peek(element -> PyUtil.verboseOnly(() -> PyPsiUtils.assertValid(element))) + .map(element -> getCallTargetReturnType(call, element, context)) + .nonNull() + .map(Ref::get) + .toList(); + if (!members.isEmpty()) { return PyUnionType.union(members); } @@ -839,27 +842,29 @@ public class PyCallExpressionHelper { } @NotNull - public static Stream forEveryScopeTakeOverloadsOtherwiseImplementations(@NotNull List callees, - @NotNull TypeEvalContext context) { - if (!containsOverloadsAndImplementations(callees, context)) { - return callees.stream(); + public static Stream forEveryScopeTakeOverloadsOtherwiseImplementations(@NotNull List elements, + @NotNull Function mapper, + @NotNull TypeEvalContext context) { + if (!containsOverloadsAndImplementations(elements, mapper, context)) { + return elements.stream(); } return StreamEx - .of(callees) - .groupingBy(callee -> ScopeUtil.getScopeOwner(callee.getElement())) + .of(elements) + .groupingBy(element -> ScopeUtil.getScopeOwner(mapper.apply(element))) .values() .stream() - .flatMap(oneScopeCallees -> takeOverloadsOtherwiseImplementations(oneScopeCallees, context)); + .flatMap(oneScopeElements -> takeOverloadsOtherwiseImplementations(oneScopeElements, mapper, context)); } - private static boolean containsOverloadsAndImplementations(@NotNull List callees, - @NotNull TypeEvalContext context) { + private static boolean containsOverloadsAndImplementations(@NotNull List elements, + @NotNull Function mapper, + @NotNull TypeEvalContext context) { boolean containsOverloads = false; boolean containsImplementations = false; - for (PyCallExpression.PyRatedMarkedCallee callee : callees) { - final boolean overload = PyiUtil.isOverload(callee.getElement(), context); + for (E element : elements) { + final boolean overload = PyiUtil.isOverload(mapper.apply(element), context); containsOverloads |= overload; containsImplementations |= !overload; @@ -870,13 +875,14 @@ public class PyCallExpressionHelper { } @NotNull - private static Stream takeOverloadsOtherwiseImplementations(@NotNull List callees, - @NotNull TypeEvalContext context) { - if (!containsOverloadsAndImplementations(callees, context)) { - return callees.stream(); + private static Stream takeOverloadsOtherwiseImplementations(@NotNull List elements, + @NotNull Function mapper, + @NotNull TypeEvalContext context) { + if (!containsOverloadsAndImplementations(elements, mapper, context)) { + return elements.stream(); } - return callees.stream().filter(callee -> PyiUtil.isOverload(callee.getElement(), context)); + return elements.stream().filter(element -> PyiUtil.isOverload(mapper.apply(element), context)); } public static class ArgumentMappingResults { diff --git a/python/src/com/jetbrains/python/psi/types/PyTypeChecker.java b/python/src/com/jetbrains/python/psi/types/PyTypeChecker.java index 6ae5740572d2..e165c112b574 100644 --- a/python/src/com/jetbrains/python/psi/types/PyTypeChecker.java +++ b/python/src/com/jetbrains/python/psi/types/PyTypeChecker.java @@ -18,6 +18,7 @@ package com.jetbrains.python.psi.types; import com.intellij.openapi.extensions.Extensions; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiNamedElement; +import com.intellij.psi.ResolveResult; import com.intellij.util.ArrayUtil; import com.intellij.util.containers.ContainerUtil; import com.jetbrains.python.PyNames; @@ -627,7 +628,7 @@ public class PyTypeChecker { final List ratedMarkedCallees = PyUtil.filterTopPriorityResults(((PyCallExpression)callSite).multiResolveRatedCallee(resolveContext)); - return forEveryScopeTakeOverloadsOtherwiseImplementations(ratedMarkedCallees, context) + return forEveryScopeTakeOverloadsOtherwiseImplementations(ratedMarkedCallees, ResolveResult::getElement, context) .map(PyCallExpression.PyRatedMarkedCallee::getElement) .collect(Collectors.toList()); } diff --git a/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java b/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java index ddde12e6cb3c..820b8a4a1331 100644 --- a/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java +++ b/python/src/com/jetbrains/python/pyi/PyiTypeProvider.java @@ -91,9 +91,8 @@ public class PyiTypeProvider extends PyTypeProviderBase { if (pythonStub instanceof PyFunction) { return getOverloadedCallType((PyFunction)pythonStub, callSite, context); } - else if (function.getContainingFile() instanceof PyiFile) { - return getOverloadedCallType(function, callSite, context); - } + + return getOverloadedCallType(function, callSite, context); } return null; diff --git a/python/testData/types/FirstOverloadAndImplementationInImportedClass/b/__init__.py b/python/testData/types/FirstOverloadAndImplementationInImportedClass/b/__init__.py new file mode 100644 index 000000000000..a2e9af286330 --- /dev/null +++ b/python/testData/types/FirstOverloadAndImplementationInImportedClass/b/__init__.py @@ -0,0 +1,14 @@ +from typing import overload + + +class A: + @overload + def foo(self, value: int) -> int: + pass + + @overload + def foo(self, value: str) -> str: + pass + + def foo(self, value): + return None \ No newline at end of file diff --git a/python/testData/types/FirstOverloadAndImplementationInImportedModule/b/__init__.py b/python/testData/types/FirstOverloadAndImplementationInImportedModule/b/__init__.py new file mode 100644 index 000000000000..6cbc514c2e2e --- /dev/null +++ b/python/testData/types/FirstOverloadAndImplementationInImportedModule/b/__init__.py @@ -0,0 +1,15 @@ +from typing import overload + + +@overload +def foo(value: int) -> int: + pass + + +@overload +def foo(value: str) -> str: + pass + + +def foo(value): + return None \ No newline at end of file diff --git a/python/testData/types/NotMatchedOverloadsAndImplementationInImportedClass/b/__init__.py b/python/testData/types/NotMatchedOverloadsAndImplementationInImportedClass/b/__init__.py new file mode 100644 index 000000000000..a2e9af286330 --- /dev/null +++ b/python/testData/types/NotMatchedOverloadsAndImplementationInImportedClass/b/__init__.py @@ -0,0 +1,14 @@ +from typing import overload + + +class A: + @overload + def foo(self, value: int) -> int: + pass + + @overload + def foo(self, value: str) -> str: + pass + + def foo(self, value): + return None \ No newline at end of file diff --git a/python/testData/types/NotMatchedOverloadsAndImplementationInImportedModule/b/__init__.py b/python/testData/types/NotMatchedOverloadsAndImplementationInImportedModule/b/__init__.py new file mode 100644 index 000000000000..6cbc514c2e2e --- /dev/null +++ b/python/testData/types/NotMatchedOverloadsAndImplementationInImportedModule/b/__init__.py @@ -0,0 +1,15 @@ +from typing import overload + + +@overload +def foo(value: int) -> int: + pass + + +@overload +def foo(value: str) -> str: + pass + + +def foo(value): + return None \ No newline at end of file diff --git a/python/testData/types/SecondOverloadAndImplementationInImportedClass/b/__init__.py b/python/testData/types/SecondOverloadAndImplementationInImportedClass/b/__init__.py new file mode 100644 index 000000000000..a2e9af286330 --- /dev/null +++ b/python/testData/types/SecondOverloadAndImplementationInImportedClass/b/__init__.py @@ -0,0 +1,14 @@ +from typing import overload + + +class A: + @overload + def foo(self, value: int) -> int: + pass + + @overload + def foo(self, value: str) -> str: + pass + + def foo(self, value): + return None \ No newline at end of file diff --git a/python/testData/types/SecondOverloadAndImplementationInImportedModule/b/__init__.py b/python/testData/types/SecondOverloadAndImplementationInImportedModule/b/__init__.py new file mode 100644 index 000000000000..6cbc514c2e2e --- /dev/null +++ b/python/testData/types/SecondOverloadAndImplementationInImportedModule/b/__init__.py @@ -0,0 +1,15 @@ +from typing import overload + + +@overload +def foo(value: int) -> int: + pass + + +@overload +def foo(value: str) -> str: + pass + + +def foo(value): + return None \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java index 4eb2a26cce7f..85fa183469b9 100644 --- a/python/testSrc/com/jetbrains/python/PyTypeTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java @@ -1714,6 +1714,177 @@ public class PyTypeTest extends PyTestCase { getTypeEvalContexts(expression).forEach(context -> context.getType(expression)); } + // PY-22971 + public void testFirstOverloadAndImplementationInClass() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doTest("int", + "from typing import overload\n" + + "class A:\n" + + " @overload\n" + + " def foo(self, value: int) -> int:\n" + + " pass\n" + + " @overload\n" + + " def foo(self, value: str) -> str:\n" + + " pass\n" + + " def foo(self, value):\n" + + " return None\n" + + "expr = A().foo(5)") + ); + } + + // PY-22971 + public void testTopLevelFirstOverloadAndImplementation() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doTest("int", + "from typing import overload\n" + + "@overload\n" + + "def foo(value: int) -> int:\n" + + " pass\n" + + "@overload\n" + + "def foo(value: str) -> str:\n" + + " pass\n" + + "def foo(value):\n" + + " return None\n" + + "expr = foo(5)") + ); + } + + // PY-22971 + public void testFirstOverloadAndImplementationInImportedClass() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doMultiFileTest("int", + "from b import A\n" + + "expr = A().foo(5)") + ); + } + + // PY-22971 + public void testFirstOverloadAndImplementationInImportedModule() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doMultiFileTest("int", + "from b import foo\n" + + "expr = foo(5)") + ); + } + + // PY-22971 + public void testSecondOverloadAndImplementationInClass() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doTest("str", + "from typing import overload\n" + + "class A:\n" + + " @overload\n" + + " def foo(self, value: int) -> int:\n" + + " pass\n" + + " @overload\n" + + " def foo(self, value: str) -> str:\n" + + " pass\n" + + " def foo(self, value):\n" + + " return None\n" + + "expr = A().foo(\"5\")") + ); + } + + // PY-22971 + public void testTopLevelSecondOverloadAndImplementation() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doTest("str", + "from typing import overload\n" + + "@overload\n" + + "def foo(value: int) -> int:\n" + + " pass\n" + + "@overload\n" + + "def foo(value: str) -> str:\n" + + " pass\n" + + "def foo(value):\n" + + " return None\n" + + "expr = foo(\"5\")") + ); + } + + // PY-22971 + public void testSecondOverloadAndImplementationInImportedClass() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doMultiFileTest("str", + "from b import A\n" + + "expr = A().foo(\"5\")") + ); + } + + // PY-22971 + public void testSecondOverloadAndImplementationInImportedModule() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doMultiFileTest("str", + "from b import foo\n" + + "expr = foo(\"5\")") + ); + } + + // PY-22971 + public void testNotMatchedOverloadsAndImplementationInClass() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doTest("Union[int, str]", + "from typing import overload\n" + + "class A:\n" + + " @overload\n" + + " def foo(self, value: int) -> int:\n" + + " pass\n" + + " @overload\n" + + " def foo(self, value: str) -> str:\n" + + " pass\n" + + " def foo(self, value):\n" + + " return None\n" + + "expr = A().foo(object())") + ); + } + + // PY-22971 + public void testTopLevelNotMatchedOverloadsAndImplementation() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doTest("Union[int, str]", + "from typing import overload\n" + + "@overload\n" + + "def foo(value: int) -> int:\n" + + " pass\n" + + "@overload\n" + + "def foo(value: str) -> str:\n" + + " pass\n" + + "def foo(value):\n" + + " return None\n" + + "expr = foo(object())") + ); + } + + // PY-22971 + public void testNotMatchedOverloadsAndImplementationInImportedClass() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doMultiFileTest("Union[int, str]", + "from b import A\n" + + "expr = A().foo(object())") + ); + } + + // PY-22971 + public void testNotMatchedOverloadsAndImplementationInImportedModule() { + runWithLanguageLevel( + LanguageLevel.PYTHON35, + () -> doMultiFileTest("Union[int, str]", + "from b import foo\n" + + "expr = foo(object())") + ); + } + private static List getTypeEvalContexts(@NotNull PyExpression element) { return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(), TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());