PY-46356 Suggest qualified import names for common aliases

I moved the custom index lookup, previously used in PyImportCollector to find
modules and packages, to PyModuleNameIndex, already utilized in other similar,
places such as PyModulePackageCompletionContributor. Two new methods where
introduced to the interface of PyModuleNameIndex for that. The first one
is a generalization of  find(), allowing to re-use an existing search scope, and
the second finds modules by their fully qualified name in the same fashion as
it was done in PyClassNameIndex.findClass().

GitOrigin-RevId: 2de5821351eacf08ddc045454eea189b21fa1186
This commit is contained in:
Mikhail Golubev
2021-01-07 12:17:33 +03:00
committed by intellij-monorepo-bot
parent 8e53b39d77
commit 6f5997ee45
12 changed files with 89 additions and 34 deletions

View File

@@ -121,7 +121,7 @@ public class ImportFromExistingAction implements QuestionAction {
if (manager.isInjectedFragment(file)) {
file = manager.getTopLevelFile(myTarget);
}
// We are trying to import top-level module or package which thus cannot be qualified
// A root-level module or package cannot be imported with a "from" import.
if (PyUtil.isRoot(item.getFile())) {
if (myImportLocally) {
AddImportHelper.addLocalImportStatement(myTarget, item.getImportableName());
@@ -143,7 +143,9 @@ public class ImportFromExistingAction implements QuestionAction {
else {
AddImportHelper.addImportStatement(file, nameToImport, item.getAsName(), priority, myTarget);
}
myTarget.replace(gen.createExpressionFromText(LanguageLevel.forElement(myTarget), qualifiedName + "." + myName));
if (item.getAsName() == null) {
myTarget.replace(gen.createExpressionFromText(LanguageLevel.forElement(myTarget), qualifiedName + "." + myName));
}
}
else {
if (myImportLocally) {

View File

@@ -5,7 +5,6 @@ import com.intellij.openapi.progress.ProgressManager;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.io.FileUtilRt;
import com.intellij.psi.*;
import com.intellij.psi.search.FilenameIndex;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.psi.util.QualifiedName;
@@ -18,7 +17,9 @@ import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
import com.jetbrains.python.psi.search.PySearchUtilBase;
import com.jetbrains.python.psi.stubs.PyClassNameIndex;
import com.jetbrains.python.psi.stubs.PyFunctionNameIndex;
import com.jetbrains.python.psi.stubs.PyModuleNameIndex;
import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
import org.jetbrains.annotations.NotNull;
import java.util.*;
@@ -115,10 +116,10 @@ public class PyImportCollector {
}
symbols.addAll(PyVariableNameIndex.find(myRefText, project, scope));
if (isPossibleModuleReference()) {
symbols.addAll(findImportableModules(myRefText, false, project, scope));
symbols.addAll(findImportableModules(myRefText, false, scope));
String packageQName = PyPackageAliasesProvider.commonImportAliases.get(myRefText);
if (packageQName != null) {
symbols.addAll(findImportableModules(packageQName, true, project, scope));
symbols.addAll(findImportableModules(packageQName, true, scope));
}
}
for (PsiNamedElement symbol : symbols) {
@@ -190,36 +191,20 @@ public class PyImportCollector {
return true;
}
private Collection<PsiFileSystemItem> findImportableModules(String name,
@NotNull
private Collection<PsiFileSystemItem> findImportableModules(@NotNull String name,
boolean matchQualifiedName,
Project project,
GlobalSearchScope scope) {
@NotNull GlobalSearchScope scope) {
List<PsiFileSystemItem> result = new ArrayList<>();
// Add packages
QualifiedName qualifiedName = QualifiedName.fromDottedString(name);
FilenameIndex.processFilesByName(name, true, item -> {
ProgressManager.checkCanceled();
final PsiDirectory candidatePackageDir = as(item, PsiDirectory.class);
if (candidatePackageDir != null && candidatePackageDir.findFile(PyNames.INIT_DOT_PY) != null) {
QualifiedName shortestName = QualifiedNameFinder.findShortestImportableQName(candidatePackageDir);
if (!matchQualifiedName || qualifiedName.equals(shortestName)) {
result.add(candidatePackageDir);
}
List<PyFile> matchingModules = matchQualifiedName ? PyModuleNameIndex.findByQualifiedName(qualifiedName, myNode.getProject(), scope)
: PyModuleNameIndex.findByShortName(name, myNode.getProject(), scope);
for (PyFile module : matchingModules) {
PsiFileSystemItem candidate = as(PyUtil.turnInitIntoDir(module), PsiFileSystemItem.class);
if (candidate != null && PyUtil.isImportable(myNode.getContainingFile(), candidate)) {
result.add(candidate);
}
return true;
}, scope, project, null);
// Add modules
FilenameIndex.processFilesByName(name + ".py", false, true, item -> {
ProgressManager.checkCanceled();
if (PyUtil.isImportable(myNode.getContainingFile(), item)) {
QualifiedName shortestName = QualifiedNameFinder.findShortestImportableQName(item);
if (!matchQualifiedName || qualifiedName.equals(shortestName)) {
result.add(item);
}
}
return true;
}, scope, project, null);
}
return result;
}

View File

@@ -5,9 +5,12 @@ import com.intellij.openapi.fileTypes.FileTypeRegistry;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.io.FileUtilRt;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiManager;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.containers.ContainerUtil;
import com.intellij.util.indexing.*;
import com.intellij.util.io.EnumeratorStringDescriptor;
import com.intellij.util.io.KeyDescriptor;
@@ -15,6 +18,7 @@ import com.jetbrains.python.PyNames;
import com.jetbrains.python.PythonFileType;
import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
import com.jetbrains.python.psi.PyFile;
import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
import com.jetbrains.python.psi.search.PySearchUtilBase;
import org.jetbrains.annotations.NotNull;
@@ -84,14 +88,28 @@ public class PyModuleNameIndex extends ScalarIndexExtension<String> {
}
@NotNull
public static List<PyFile> find(@NotNull String name, @NotNull Project project, boolean includeNonProjectItems) {
final List<PyFile> results = new ArrayList<>();
public static List<PyFile> find(@NotNull String shortName, @NotNull Project project, boolean includeNonProjectItems) {
final GlobalSearchScope baseScope = includeNonProjectItems
? PySearchUtilBase.excludeSdkTestsScope(project)
: GlobalSearchScope.projectScope(project);
final GlobalSearchScope scope = baseScope
.intersectWith(GlobalSearchScope.notScope(PyUserSkeletonsUtil.getUserSkeletonsDirectoryScope(project)));
final Collection<VirtualFile> files = FileBasedIndex.getInstance().getContainingFiles(NAME, name, scope);
return findByShortName(shortName, project, scope);
}
/**
* Returns all modules with the given short name (the last component of a fully qualified name).
* <p>
* File extensions should not be included. For __init__.py modules, the name of the corresponding directory is matched.
*
* @param shortName short name of a module or name of the containing package for __init__.py modules
* @param project project where the search is performed
* @param scope search scope, limiting applicable virtual files
*/
@NotNull
public static List<PyFile> findByShortName(@NotNull String shortName, @NotNull Project project, @NotNull GlobalSearchScope scope) {
final List<PyFile> results = new ArrayList<>();
final Collection<VirtualFile> files = FileBasedIndex.getInstance().getContainingFiles(NAME, shortName, scope);
for (VirtualFile virtualFile : files) {
final PsiFile psiFile = PsiManager.getInstance(project).findFile(virtualFile);
if (psiFile instanceof PyFile) {
@@ -100,4 +118,27 @@ public class PyModuleNameIndex extends ScalarIndexExtension<String> {
}
return results;
}
/**
* Returns all modules with the given fully qualified name.
* <p>
* For __init__.py modules, the qualified name of the corresponding package is used.
* <p>
* All possible qualified names of a module are considered. For instance, in case of a source root "src" inside a project's
* content root, src/foo.py or src/foo/__init__.py will be returned both for qualified names "foo" and "src.foo".
*
* @param qName short name of a module or name of the containing package for __init__.py modules
* @param project project where the search is performed
* @param scope search scope, limiting applicable virtual files
* @see QualifiedNameFinder#findImportableQNames(PsiElement, VirtualFile)
*/
@NotNull
public static List<PyFile> findByQualifiedName(@NotNull QualifiedName qName, @NotNull Project project, @NotNull GlobalSearchScope scope) {
String shortName = qName.getLastComponent();
if (shortName == null) return Collections.emptyList();
return ContainerUtil.mapNotNull(findByShortName(shortName, project, scope), file -> {
List<QualifiedName> possibleQNames = QualifiedNameFinder.findImportableQNames(file, file.getVirtualFile());
return possibleQNames.contains(qName) ? file : null;
});
}
}

View File

@@ -0,0 +1 @@
<error descr="Unresolved reference 'plt'">p<caret>lt</error>.plot()

View File

@@ -0,0 +1,3 @@
from matplotlib import pyplot as plt
plt.plot()

View File

@@ -0,0 +1 @@
<error descr="Unresolved reference 'plt'">p<caret>lt</error>.plot()

View File

@@ -0,0 +1,3 @@
import matplotlib.pyplot as plt
plt.plot()

View File

@@ -25,6 +25,7 @@ import com.intellij.util.ObjectUtils;
import com.intellij.util.Processor;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyQuickFixTestCase;
import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
import com.jetbrains.python.codeInsight.imports.AutoImportQuickFix;
import com.jetbrains.python.codeInsight.imports.ImportCandidateHolder;
import com.jetbrains.python.codeInsight.imports.PythonImportUtils;
@@ -297,6 +298,24 @@ public class PyAddImportQuickFixTest extends PyQuickFixTestCase {
doMultiFileAutoImportTest("Import");
}
// PY-46356
public void testCommonSubModuleAliasPlainImport() {
PyCodeInsightSettings codeInsightSettings = PyCodeInsightSettings.getInstance();
boolean oldPreferFromImport = codeInsightSettings.PREFER_FROM_IMPORT;
codeInsightSettings.PREFER_FROM_IMPORT = false;
try {
doMultiFileAutoImportTest("Qualify with an imported module");
}
finally {
codeInsightSettings.PREFER_FROM_IMPORT = oldPreferFromImport;
}
}
// PY-46356
public void testCommonSubModuleAliasFromImport() {
doMultiFileAutoImportTest("Import");
}
private void doTestProposedImportsOrdering(String @NotNull ... expected) {
doMultiFileAutoImportTest("Import", fix -> {
final List<String> candidates = ContainerUtil.map(fix.getCandidates(), c -> c.getPresentableText());