Multi-resolve types in docstrings (PY-18427)

This commit is contained in:
Andrey Vlasovskikh
2016-02-08 01:05:13 +03:00
parent 34da6faf01
commit fb26f5a9a4
2 changed files with 48 additions and 26 deletions
@@ -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<PyElementType> 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<RatedResolveResult> 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<PyType> 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
@@ -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<TypeEvalContext> getTypeEvalContexts(@NotNull PyExpression element) {
return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(),
TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());