Multi-resolve types in PEP 484 type hints (PY-18427)

This commit is contained in:
Andrey Vlasovskikh
2016-02-08 01:05:15 +03:00
parent fb26f5a9a4
commit 795d640d40
2 changed files with 67 additions and 39 deletions
@@ -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<PyType> 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<String> 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<String> 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<String> 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<String> 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<PsiElement> tryResolving(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
final List<PsiElement> 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<PsiElement> 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.<PsiElement>singletonList(expression);
}
@Nullable
private static String resolveToQualifiedName(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
return getQualifiedName(tryResolving(expression, context));
@NotNull
private static Collection<String> resolveToQualifiedNames(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
final Set<String> names = Sets.newLinkedHashSet();
for (PsiElement resolved : tryResolving(expression, context)) {
final String name = getQualifiedName(resolved);
if (name != null) {
names.add(name);
}
}
return names;
}
@Nullable
@@ -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());