diff --git a/python/src/com/jetbrains/python/psi/types/PyTypeParser.java b/python/src/com/jetbrains/python/psi/types/PyTypeParser.java index d721a9051658..bbd5f6675752 100644 --- a/python/src/com/jetbrains/python/psi/types/PyTypeParser.java +++ b/python/src/com/jetbrains/python/psi/types/PyTypeParser.java @@ -16,6 +16,7 @@ package com.jetbrains.python.psi.types; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Lists; import com.intellij.openapi.util.Pair; import com.intellij.openapi.util.TextRange; import com.intellij.psi.PsiDirectory; @@ -438,42 +439,48 @@ public class PyTypeParser { final Token firstToken = tokens.get(0); final String firstText = firstToken.getText().toString(); final TextRange firstRange = firstToken.getRange(); - final PsiElement resolved = file.getElementNamed(firstText); - if (resolved == null) { + final List resolveResults = file.multiResolveName(firstText); + if (resolveResults.isEmpty()) { return getImplicitlyResolvedType(tokens, context, types, fullRanges, firstRange); } - PyType type = null; - // Local or imported name - if (resolved instanceof PyTargetExpression) { - type = PyTypingTypeProvider.getTypeFromTargetExpression((PyTargetExpression)resolved, context); - } - if (type == null && resolved instanceof PyTypedElement) { - type = context.getType((PyTypedElement)resolved); - } - if (type != null) { - tokens.remove(0); - if (!allowResolveToType(type)) { - return null; + final List members = Lists.newArrayList(); + for (RatedResolveResult result : resolveResults) { + final PsiElement resolved = result.getElement(); + PyType type = null; + if (resolved instanceof PyTargetExpression) { + type = PyTypingTypeProvider.getTypeFromTargetExpression((PyTargetExpression)resolved, context); } - if (type instanceof PyClassLikeType) { - type = ((PyClassLikeType)type).toInstance(); + if (type == null && resolved instanceof PyTypedElement) { + type = context.getType((PyTypedElement)resolved); } - types.put(firstRange, type); - fullRanges.put(type, firstRange); - for (PyFromImportStatement fromImportStatement : file.getFromImports()) { - for (PyImportElement importElement : fromImportStatement.getImportElements()) { + if (type != null) { + if (!allowResolveToType(type)) { + continue; + } + if (type instanceof PyClassLikeType) { + type = ((PyClassLikeType)type).toInstance(); + } + types.put(firstRange, type); + fullRanges.put(type, firstRange); + for (PyFromImportStatement fromImportStatement : file.getFromImports()) { + for (PyImportElement importElement : fromImportStatement.getImportElements()) { + if (firstText.equals(importElement.getVisibleName())) { + imports.put(type, importElement); + } + } + } + for (PyImportElement importElement : file.getImportTargets()) { if (firstText.equals(importElement.getVisibleName())) { imports.put(type, importElement); } } } - for (PyImportElement importElement : file.getImportTargets()) { - if (firstText.equals(importElement.getVisibleName())) { - imports.put(type, importElement); - } - } + members.add(type); } - return type; + if (!members.isEmpty()) { + tokens.remove(0); + } + return PyUnionType.union(members); } @Nullable diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java index 13039998e0e4..590685ed2a99 100644 --- a/python/testSrc/com/jetbrains/python/PyTypeTest.java +++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java @@ -1025,6 +1025,21 @@ public class PyTypeTest extends PyTestCase { " expr = foo\n"); } + // PY-18427 + public void testConditionalTypeInDocstring() { + doTest("Union[str, int]", + "if something:\n" + + " Type = int\n" + + "else:\n" + + " Type = str\n" + + "\n" + + "def f(expr):\n" + + " '''\n" + + " :type expr: Type\n" + + " '''\n" + + " pass\n"); + } + private static List getTypeEvalContexts(@NotNull PyExpression element) { return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(), TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());