com.jetbrains.python.codeInsight.imports.PythonReferenceImporter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of python-community Show documentation
A packaging of the IntelliJ Community Edition python-community library. This is release number 1 of trunk branch 142.
The newest version!
/*
 * Copyright 2000-2014 JetBrains s.r.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * User: anna
 * Date: 11-Mar-2008
 */
package com.jetbrains.python.codeInsight.imports;

import com.intellij.codeInsight.daemon.ReferenceImporter;
import com.intellij.codeInsight.daemon.impl.CollectHighlightsUtil;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.extensions.Extensions;
import com.intellij.openapi.module.Module;
import com.intellij.openapi.module.ModuleUtilCore;
import com.intellij.openapi.progress.ProgressManager;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.io.FileUtil;
import com.intellij.openapi.util.registry.Registry;
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;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
import com.jetbrains.python.inspections.unresolvedReference.PyPackageAliasesProvider;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyFileImpl;
import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
import com.jetbrains.python.psi.search.PyProjectScopeBuilder;
import com.jetbrains.python.psi.stubs.PyClassNameIndex;
import com.jetbrains.python.psi.stubs.PyFunctionNameIndex;
import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
import com.jetbrains.python.sdk.PythonSdkType;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

public class PythonReferenceImporter implements ReferenceImporter {
  @Override
  public boolean autoImportReferenceAtCursor(@NotNull final Editor editor, @NotNull final PsiFile file) {
    if (!(file instanceof PyFile)) {
      return false;
    }
    int caretOffset = editor.getCaretModel().getOffset();
    Document document = editor.getDocument();
    int lineNumber = document.getLineNumber(caretOffset);
    int startOffset = document.getLineStartOffset(lineNumber);
    int endOffset = document.getLineEndOffset(lineNumber);

    List elements = CollectHighlightsUtil.getElementsInRange(file, startOffset, endOffset);
    for (PsiElement element : elements) {
      if (element instanceof PyReferenceExpression && isImportable(element)) {
        final PyReferenceExpression refExpr = (PyReferenceExpression)element;
        if (!refExpr.isQualified()) {
          final PsiPolyVariantReference reference = refExpr.getReference();
          if (reference.resolve() == null) {
            AutoImportQuickFix fix = proposeImportFix(refExpr, reference);
            if (fix != null && fix.getCandidatesCount() == 1) {
              fix.invoke(file);
            }
            return true;
          }
        }
      }
    }
    return false;
  }

  @Override
  public boolean autoImportReferenceAt(@NotNull Editor editor, @NotNull PsiFile file, int offset) {
    if (!(file instanceof PyFile)) {
      return false;
    }
    PsiReference element = file.findReferenceAt(offset);
    if (element instanceof PyReferenceExpression && isImportable((PsiElement)element)) {
      final PyReferenceExpression refExpr = (PyReferenceExpression)element;
      if (!refExpr.isQualified()) {
        final PsiPolyVariantReference reference = refExpr.getReference();
        if (reference.resolve() == null) {
          AutoImportQuickFix fix = proposeImportFix(refExpr, reference);
          if (fix != null && fix.getCandidatesCount() == 1) {
            fix.invoke(file);
          }
          return true;
        }
      }
    }
    return false;
  }

  @Nullable
  public static AutoImportQuickFix proposeImportFix(final PyElement node, PsiReference reference) {
    final String text = reference.getElement().getText();
    final String refText = reference.getRangeInElement().substring(text); // text of the part we're working with

    // don't propose meaningless auto imports if no interpreter is configured
    final Module module = ModuleUtilCore.findModuleForPsiElement(node);
    if (module != null && PythonSdkType.findPythonSdk(module) == null) {
      return null;
    }

    // don't show auto-import fix if we're trying to reference a variable which is defined below in the same scope
    ScopeOwner scopeOwner = PsiTreeUtil.getParentOfType(node, ScopeOwner.class);
    if (scopeOwner != null && ControlFlowCache.getScope(scopeOwner).containsDeclaration(refText)) {
      return null;
    }

    AutoImportQuickFix fix = addCandidates(node, reference, refText, null);
    if (fix != null) return fix;
    final String packageName = PyPackageAliasesProvider.commonImportAliases.get(refText);
    if (packageName != null) {
      fix = addCandidates(node, reference, packageName, refText);
      if (fix != null) return fix;
    }
    return null;
  }

  @Nullable
  private static AutoImportQuickFix addCandidates(PyElement node, PsiReference reference, String refText, @Nullable String asName) {
    AutoImportQuickFix fix = new AutoImportQuickFix(node, reference, refText, !PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT);
    Set seenFileNames = new HashSet(); // true import names

    PsiFile existingImportFile = addCandidatesFromExistingImports(node, refText, fix, seenFileNames);
    if (fix.getCandidatesCount() == 0 || fix.hasProjectImports() || Registry.is("python.import.always.ask")) {
      // maybe some unimported file has it, too
      ProgressManager.checkCanceled(); // before expensive index searches
      addSymbolImportCandidates(node, refText, asName, fix, seenFileNames, existingImportFile);
    }

    for(PyImportCandidateProvider provider: Extensions.getExtensions(PyImportCandidateProvider.EP_NAME)) {
      provider.addImportCandidates(reference, refText, fix);
    }
    if (fix.getCandidatesCount() > 0) {
      fix.sortCandidates();
      return fix;
    }
    return null;
  }

  /**
   * maybe the name is importable via some existing 'import foo' statement, and only needs a qualifier.
   * collect all such statements and analyze.
   * NOTE: It only makes sense to look at imports in file scope - there is no guarantee that an import in a local scope will
   * be visible from the scope where the auto-import was invoked
   *
   * @param node
   * @param refText
   * @param fix
   * @param seenFileNames
   * @return
   */
  @Nullable
  private static PsiFile addCandidatesFromExistingImports(PyElement node, String refText, AutoImportQuickFix fix,
                                                          Set seenFileNames) {
    PsiFile existingImportFile = null; // if there's a matching existing import, this it the file it imports
    PsiFile file = node.getContainingFile();
    if (file instanceof PyFile) {
      PyFile pyFile = (PyFile)file;
      for (PyImportElement importElement : pyFile.getImportTargets()) {
        existingImportFile = addImportViaElement(refText, fix, seenFileNames, existingImportFile, importElement, importElement.resolve());
      }
      for (PyFromImportStatement fromImportStatement : pyFile.getFromImports()) {
        if (!(fromImportStatement.isStarImport()) && fromImportStatement.getImportElements().length > 0) {
          PsiElement source = fromImportStatement.resolveImportSource();
          existingImportFile = addImportViaElement(refText, fix, seenFileNames, existingImportFile, fromImportStatement.getImportElements()[0], source);
        }
      }
    }
    return existingImportFile;
  }

  private static PsiFile addImportViaElement(String refText,
                                             AutoImportQuickFix fix,
                                             Set seenFileNames,
                                             PsiFile existingImportFile,
                                             PyImportElement importElement,
                                             PsiElement source) {
    PsiElement sourceFile = PyUtil.turnDirIntoInit(source);
    if (sourceFile instanceof PyFileImpl) {
      PyStatement importStatement = importElement.getContainingImportStatement();
      String refName = null;
      if (importStatement instanceof PyFromImportStatement) {
        QualifiedName qName = ((PyFromImportStatement)importStatement).getImportSourceQName();
        if (qName != null) {
          refName = qName.toString();
        }
      }
      else {
        QualifiedName importReferenceQName = importElement.getImportedQName();
        if (importReferenceQName != null) {
          refName = importReferenceQName.toString();
        }
      }
      if (refName != null) {
        if (seenFileNames.contains(refName)) {
          return existingImportFile;
        }
        seenFileNames.add(refName);
      }

      PyFileImpl importSourceFile = (PyFileImpl)sourceFile;
      PsiElement res = importSourceFile.findExportedName(refText);
      // allow importing from this source if it either declares the name itself or represents a higher-level package that reexports the name
      if (res != null && !(res instanceof PyFile) && !(res instanceof PyImportElement) && res.getContainingFile() != null &&
          PsiTreeUtil.isAncestor(source, res.getContainingFile(), false)) {
        existingImportFile = importSourceFile;
        fix.addImport(res, importSourceFile, importElement);
      }
    }
    return existingImportFile;
  }

  private static void addSymbolImportCandidates(PyElement node, String refText, @Nullable String asName, AutoImportQuickFix fix,
                                                Set seenFileNames, PsiFile existingImportFile) {
    Project project = node.getProject();
    List symbols = new ArrayList();
    symbols.addAll(PyClassNameIndex.find(refText, project, true));
    GlobalSearchScope scope = PyProjectScopeBuilder.excludeSdkTestsScope(node);
    if (!isQualifier(node)) {
      symbols.addAll(PyFunctionNameIndex.find(refText, project, scope));
    }
    symbols.addAll(PyVariableNameIndex.find(refText, project, scope));
    if (isPossibleModuleReference(node)) {
      symbols.addAll(findImportableModules(node.getContainingFile(), refText, project, scope));
    }
    if (!symbols.isEmpty()) {
      for (PsiElement symbol : symbols) {
        if (isIndexableTopLevel(symbol)) { // we only want top-level symbols
          PsiFileSystemItem srcfile = symbol instanceof PsiFileSystemItem ? ((PsiFileSystemItem)symbol).getParent() : symbol.getContainingFile();
          if (srcfile != null && isAcceptableForImport(node, existingImportFile, srcfile)) {
            QualifiedName importPath = QualifiedNameFinder.findCanonicalImportPath(symbol, node);
            if (symbol instanceof PsiFileSystemItem && importPath != null) {
              importPath = importPath.removeTail(1);
            }
            if (importPath != null && !seenFileNames.contains(importPath.toString())) {
              // a new, valid hit
              fix.addImport(symbol, srcfile, importPath, asName);
              seenFileNames.add(importPath.toString()); // just in case, again
            }
          }
        }
      }
    }
  }

  private static boolean isAcceptableForImport(PyElement node, PsiFile existingImportFile, PsiFileSystemItem srcfile) {
    return srcfile != existingImportFile && srcfile != node.getContainingFile() &&
        (ImportFromExistingAction.isRoot(srcfile) || PyNames.isIdentifier(FileUtil.getNameWithoutExtension(srcfile.getName()))) &&
         !isShadowedModule(srcfile);
  }

  private static boolean isShadowedModule(PsiFileSystemItem file) {
    if (file.isDirectory() || file.getName().equals(PyNames.INIT_DOT_PY)) {
      return false;
    }
    String name = FileUtil.getNameWithoutExtension(file.getName());
    final PsiDirectory directory = ((PsiFile)file).getContainingDirectory();
    if (directory == null) {
      return false;
    }
    PsiDirectory packageDir = directory.findSubdirectory(name);
    return packageDir != null && packageDir.findFile(PyNames.INIT_DOT_PY) != null;
  }

  private static boolean isQualifier(PyElement node) {
    return node.getParent() instanceof PyReferenceExpression && node == ((PyReferenceExpression)node.getParent()).getQualifier();
  }

  private static boolean isPossibleModuleReference(PyElement node) {
    if (node.getParent() instanceof PyCallExpression && node == ((PyCallExpression) node.getParent()).getCallee()) {
      return false;
    }
    if (node.getParent() instanceof PyArgumentList) {
      final PyArgumentList argumentList = (PyArgumentList)node.getParent();
      if (argumentList.getParent() instanceof PyClass) {
        final PyClass pyClass = (PyClass)argumentList.getParent();
        if (pyClass.getSuperClassExpressionList() == argumentList) {
          return false;
        }
      }
    }
    return true;
  }

  private static Collection findImportableModules(PsiFile targetFile, String reftext, Project project, GlobalSearchScope scope) {
    List result = new ArrayList();
    PsiFile[] files = FilenameIndex.getFilesByName(project, reftext + ".py", scope);
    for (PsiFile file : files) {
      if (isImportableModule(targetFile, file)) {
        result.add(file);
      }
    }
    // perhaps the module is a directory, not a file
    PsiFile[] initFiles = FilenameIndex.getFilesByName(project, PyNames.INIT_DOT_PY, scope);
    for (PsiFile initFile : initFiles) {
      PsiDirectory parent = initFile.getParent();
      if (parent != null && parent.getName().equals(reftext)) {
        result.add(parent);
      }
    }
    return result;
  }

  public static boolean isImportableModule(PsiFile targetFile, @NotNull PsiFileSystemItem file) {
    PsiDirectory parent = (PsiDirectory)file.getParent();
    return parent != null && file != targetFile &&
           (parent.findFile(PyNames.INIT_DOT_PY) != null ||
            ImportFromExistingAction.isRoot(parent) ||
            parent == targetFile.getParent());
  }

  private static boolean isIndexableTopLevel(PsiElement symbol) {
    if (symbol instanceof PsiFileSystemItem) {
      return true;
    }
    if (symbol instanceof PyClass || symbol instanceof PyFunction) {
      return PyUtil.isTopLevel(symbol);
    }
    // only top-level target expressions are included in VariableNameIndex
    return symbol instanceof PyTargetExpression;
  }

  public static boolean isImportable(PsiElement ref_element) {
    PyStatement parentStatement = PsiTreeUtil.getParentOfType(ref_element, PyStatement.class);
    if (parentStatement instanceof PyGlobalStatement || parentStatement instanceof PyNonlocalStatement ||
      parentStatement instanceof PyImportStatementBase) {
      return false;
    }
    return PsiTreeUtil.getParentOfType(ref_element, PyStringLiteralExpression.class, false, PyStatement.class) == null;
  }
}