diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java index 4c44890d7836..d1c3904ccbf5 100644 --- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java @@ -17,6 +17,8 @@ package com.jetbrains.python.codeInsight; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Sets; import com.intellij.openapi.project.Project; import com.intellij.openapi.util.Ref; import com.intellij.psi.PsiComment; @@ -31,9 +33,7 @@ import com.jetbrains.python.psi.types.*; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; +import java.util.*; import java.util.regex.Matcher; import java.util.regex.Pattern; @@ -228,10 +228,11 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { if (expr instanceof PySubscriptionExpression) { final PyExpression indexExpr = ((PySubscriptionExpression)expr).getIndexExpression(); if (indexExpr != null) { - final PsiElement resolved = tryResolving(indexExpr, context); - final PyGenericType genericType = getGenericType(resolved, context); - if (genericType != null) { - results.add(genericType); + for (PsiElement resolved : tryResolving(indexExpr, context)) { + final PyGenericType genericType = getGenericType(resolved, context); + if (genericType != null) { + results.add(genericType); + } } } } @@ -243,8 +244,11 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { @Nullable private static PyType getType(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { - final PsiElement resolved = tryResolving(expression, context); - return getTypeForResolvedElement(resolved, context); + final List members = Lists.newArrayList(); + for (PsiElement resolved : tryResolving(expression, context)) { + members.add(getTypeForResolvedElement(resolved, context)); + } + return PyUnionType.union(members); } @Nullable @@ -340,8 +344,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { if (element instanceof PySubscriptionExpression) { final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression)element; final PyExpression operand = subscriptionExpr.getOperand(); - final String operandName = resolveToQualifiedName(operand, context); - if ("typing.Optional".equals(operandName)) { + final Collection operandNames = resolveToQualifiedNames(operand, context); + if (operandNames.contains("typing.Optional")) { final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); if (indexExpr != null) { final PyType type = getType(indexExpr, context); @@ -391,8 +395,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { if (resolved instanceof PySubscriptionExpression) { final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression)resolved; final PyExpression operand = subscriptionExpr.getOperand(); - final String operandName = resolveToQualifiedName(operand, context); - if ("typing.Callable".equals(operandName)) { + final Collection operandNames = resolveToQualifiedNames(operand, context); + if (operandNames.contains("typing.Callable")) { final PyExpression indexExpr = subscriptionExpr.getIndexExpression(); if (indexExpr instanceof PyTupleExpression) { final PyTupleExpression tupleExpr = (PyTupleExpression)indexExpr; @@ -421,8 +425,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { if (element instanceof PySubscriptionExpression) { final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression)element; final PyExpression operand = subscriptionExpr.getOperand(); - final String operandName = resolveToQualifiedName(operand, context); - if ("typing.Union".equals(operandName)) { + final Collection operandNames = resolveToQualifiedNames(operand, context); + if (operandNames.contains("typing.Union")) { return PyUnionType.union(getIndexTypes(subscriptionExpr, context)); } } @@ -435,8 +439,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { final PyCallExpression assignedCall = (PyCallExpression)element; final PyExpression callee = assignedCall.getCallee(); if (callee != null) { - final String calleeQName = resolveToQualifiedName(callee, context); - if ("typing.TypeVar".equals(calleeQName)) { + final Collection calleeQNames = resolveToQualifiedNames(callee, context); + if (calleeQNames.contains("typing.TypeVar")) { final PyExpression[] arguments = assignedCall.getArguments(); if (arguments.length > 0) { final PyExpression firstArgument = arguments[0]; @@ -507,39 +511,51 @@ public class PyTypingTypeProvider extends PyTypeProviderBase { } @NotNull - private static PsiElement tryResolving(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { + private static List tryResolving(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { + final List elements = Lists.newArrayList(); if (expression instanceof PyReferenceExpression) { final PyReferenceExpression referenceExpr = (PyReferenceExpression)expression; final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context); final PsiPolyVariantReference reference = referenceExpr.getReference(resolveContext); - final PsiElement element = reference.resolve(); - if (element instanceof PyFunction) { - final PyFunction function = (PyFunction)element; - if (PyUtil.isInit(function)) { - final PyClass cls = function.getContainingClass(); - if (cls != null) { - return cls; + final List resolved = PyUtil.multiResolveTopPriority(reference); + for (PsiElement element : resolved) { + if (element instanceof PyFunction) { + final PyFunction function = (PyFunction)element; + if (PyUtil.isInit(function)) { + final PyClass cls = function.getContainingClass(); + if (cls != null) { + elements.add(cls); + continue; + } } } - } - else if (element instanceof PyTargetExpression) { - final PyTargetExpression targetExpr = (PyTargetExpression)element; - // XXX: Requires switching from stub to AST - final PyExpression assignedValue = targetExpr.findAssignedValue(); - if (assignedValue != null) { - return assignedValue; + else if (element instanceof PyTargetExpression) { + final PyTargetExpression targetExpr = (PyTargetExpression)element; + // XXX: Requires switching from stub to AST + final PyExpression assignedValue = targetExpr.findAssignedValue(); + if (assignedValue != null) { + elements.add(assignedValue); + continue; + } + } + if (element != null) { + elements.add(element); } } - if (element != null) { - return element; - } } - return expression; + return !elements.isEmpty() ? elements : Collections.singletonList(expression); } - @Nullable - private static String resolveToQualifiedName(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { - return getQualifiedName(tryResolving(expression, context)); + @NotNull + private static Collection resolveToQualifiedNames(@NotNull PyExpression expression, @NotNull TypeEvalContext context) { + final Set names = Sets.newLinkedHashSet(); + for (PsiElement resolved : tryResolving(expression, context)) { + final String name = getQualifiedName(resolved); + if (name != null) { + names.add(name); + } + } + return names; } @Nullable diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java index a9ae43a86e5e..d051cdbf6ae9 100644 --- a/python/testSrc/com/jetbrains/python/PyTypingTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java @@ -437,6 +437,18 @@ public class PyTypingTest extends PyTestCase { "expr = x.foo\n"); } + // PY-18427 + public void testConditionalType() { + doTest("Union[int, str]", + "if something:\n" + + " Type = int\n" + + "else:\n" + + " Type = str\n" + + "\n" + + "def f(expr: Type):\n" + + " pass\n"); + } + private void doTestNoInjectedText(@NotNull String text) { myFixture.configureByText(PythonFileType.INSTANCE, text); final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());