<commit_before>/*
 * 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.progress.ProgressManager;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.io.FileUtil;
import com.intellij.psi.*;
import com.intellij.psi.search.FilenameIndex;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.psi.tree.TokenSet;
import com.intellij.psi.util.PsiTreeUtil;
import com.jetbrains.python.PyElementTypes;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyQualifiedName;
import com.jetbrains.python.psi.resolve.CollectProcessor;
import com.jetbrains.python.psi.resolve.PyResolveUtil;
import com.jetbrains.python.psi.resolve.ResolveImportUtil;
import com.jetbrains.python.psi.stubs.PyClassNameIndex;
import com.jetbrains.python.psi.stubs.PyFunctionNameIndex;
import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

public class PythonReferenceImporter implements ReferenceImporter {
  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<PsiElement> elements = CollectHighlightsUtil.getElementsInRange(file, startOffset, endOffset);
    for (PsiElement element : elements) {
      if (element instanceof PyReferenceExpression && isImportable(element)) {
        final PyReferenceExpression refExpr = (PyReferenceExpression)element;
        if (refExpr.getQualifier() == null) {
          final PsiPolyVariantReference reference = refExpr.getReference();
          if (reference.resolve() == null) {
            AutoImportQuickFix fix = proposeImportFix(refExpr, reference, refExpr.getText());
            if (fix != null && fix.getCandidatesCount() == 1) {
              fix.invoke(file);
            }
            return true;
          }
        }
      }
    }
    return false;
  }

  private static TokenSet IS_IMPORT_STATEMENT = TokenSet.create(PyElementTypes.IMPORT_STATEMENT);

  @Nullable
  public static AutoImportQuickFix proposeImportFix(final PyElement node, PsiReference reference, String ref_text) {
    PsiFile existing_import_file = null; // if there's a matching existing import, this it the file it imports
    AutoImportQuickFix fix = new AutoImportQuickFix(node, reference, ref_text, !PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT);
    Set<String> seen_file_names = new HashSet<String>(); // true import names
    // maybe the name is importable via some existing 'import foo' statement, and only needs a qualifier.
    // walk up collecting all such statements and analyzing
    CollectProcessor import_prc = new CollectProcessor(IS_IMPORT_STATEMENT);
    PyResolveUtil.treeCrawlUp(import_prc, node);
    List<PsiElement> result = import_prc.getResult();
    if (result.size() > 0) {
      for (PsiElement stmt : import_prc.getResult()) {
        for (PyImportElement ielt : ((PyImportStatement)stmt).getImportElements()) {
          final PyReferenceExpression src = ielt.getImportReference();
          if (src != null) {
            PsiElement dst = src.getReference().resolve();
            if (dst instanceof PyFile) {
              PyFile dst_file = (PyFile)dst;
              String name = ielt.getImportReference().getReferencedName(); // ref is ok or matching would fail
              seen_file_names.add(name);
              PsiElement res = dst_file.findExportedName(ref_text);
              if (res != null && !(res instanceof PyFile) && !(res instanceof PyImportElement) && dst_file.equals(res.getContainingFile())) {
                existing_import_file = dst_file;
                fix.addImport(res, dst_file, ielt);
              }
            }
          }
        }
      }
    }

    // maybe some unimported file has it, too
    ProgressManager.checkCanceled(); // before expensive index searches
    // NOTE: current indices have limitations, only finding direct definitions of classes and functions.
    Project project = node.getProject();
    List<PsiElement> symbols = new ArrayList<PsiElement>();
    symbols.addAll(PyClassNameIndex.find(ref_text, project, true));
    GlobalSearchScope scope = PyClassNameIndex.projectWithLibrariesScope(project);
    if (!isQualifier(node)) {
      symbols.addAll(PyFunctionNameIndex.find(ref_text, project, scope));
    }
    symbols.addAll(PyVariableNameIndex.find(ref_text, project, scope));
    if (!isCall(node)) {
      symbols.addAll(findImportableModules(node.getContainingFile(), ref_text, project, scope));
    }
    if (symbols.size() > 0) {
      for (PsiElement symbol : symbols) {
        if (isTopLevel(symbol)) { // we only want top-level symbols
          PsiFileSystemItem srcfile = symbol instanceof PsiFileSystemItem ? ((PsiFileSystemItem)symbol).getParent() : symbol.getContainingFile();
          if (srcfile != null && srcfile != existing_import_file && srcfile != node.getContainingFile() &&
              (ImportFromExistingAction.isRoot(project, srcfile) || PyNames.isIdentifier(FileUtil.getNameWithoutExtension(srcfile.getName())))) {
            PyQualifiedName import_path = ResolveImportUtil.findCanonicalImportPath(srcfile, node);
            if (import_path != null && !seen_file_names.contains(import_path.toString())) {
              // a new, valid hit
              fix.addImport(symbol, srcfile, import_path, proposeAsName(node.getContainingFile(), ref_text, import_path));
              seen_file_names.add(import_path.toString()); // just in case, again
            }
          }
        }
      }
    }
    for(PyImportCandidateProvider provider: Extensions.getExtensions(PyImportCandidateProvider.EP_NAME)) {
      provider.addImportCandidates(reference, ref_text, fix);
    }
    if (fix.getCandidatesCount() > 0) {
      fix.sortCandidates();
      return fix;
    }
    return null;
  }

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

  private static boolean isCall(PyElement node) {
    if (node.getParent() instanceof PyCallExpression) {
      return node == ((PyCallExpression) node.getParent()).getCallee();
    }
    return false;
  }

  private static Collection<PsiElement> findImportableModules(PsiFile targetFile, String reftext, Project project, GlobalSearchScope scope) {
    List<PsiElement> result = new ArrayList<PsiElement>();
    PsiFile[] files = FilenameIndex.getFilesByName(project, reftext + ".py", scope);
    for (PsiFile file : files) {
      PsiDirectory parent = file.getParent();
      if (parent != null && (parent.findFile(PyNames.INIT_DOT_PY) != null ||
                             ImportFromExistingAction.isRoot(project, parent) ||
                             parent == targetFile.getParent())) {
        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) {<commit_msg>one more case for suggesting to import a file into itself (PY-979)<commit_after>/*
 * 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.progress.ProgressManager;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.io.FileUtil;
import com.intellij.psi.*;
import com.intellij.psi.search.FilenameIndex;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.psi.tree.TokenSet;
import com.intellij.psi.util.PsiTreeUtil;
import com.jetbrains.python.PyElementTypes;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyQualifiedName;
import com.jetbrains.python.psi.resolve.CollectProcessor;
import com.jetbrains.python.psi.resolve.PyResolveUtil;
import com.jetbrains.python.psi.resolve.ResolveImportUtil;
import com.jetbrains.python.psi.stubs.PyClassNameIndex;
import com.jetbrains.python.psi.stubs.PyFunctionNameIndex;
import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

public class PythonReferenceImporter implements ReferenceImporter {
  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<PsiElement> elements = CollectHighlightsUtil.getElementsInRange(file, startOffset, endOffset);
    for (PsiElement element : elements) {
      if (element instanceof PyReferenceExpression && isImportable(element)) {
        final PyReferenceExpression refExpr = (PyReferenceExpression)element;
        if (refExpr.getQualifier() == null) {
          final PsiPolyVariantReference reference = refExpr.getReference();
          if (reference.resolve() == null) {
            AutoImportQuickFix fix = proposeImportFix(refExpr, reference, refExpr.getText());
            if (fix != null && fix.getCandidatesCount() == 1) {
              fix.invoke(file);
            }
            return true;
          }
        }
      }
    }
    return false;
  }

  private static TokenSet IS_IMPORT_STATEMENT = TokenSet.create(PyElementTypes.IMPORT_STATEMENT);

  @Nullable
  public static AutoImportQuickFix proposeImportFix(final PyElement node, PsiReference reference, String ref_text) {
    PsiFile existing_import_file = null; // if there's a matching existing import, this it the file it imports
    AutoImportQuickFix fix = new AutoImportQuickFix(node, reference, ref_text, !PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT);
    Set<String> seen_file_names = new HashSet<String>(); // true import names
    // maybe the name is importable via some existing 'import foo' statement, and only needs a qualifier.
    // walk up collecting all such statements and analyzing
    CollectProcessor import_prc = new CollectProcessor(IS_IMPORT_STATEMENT);
    PyResolveUtil.treeCrawlUp(import_prc, node);
    List<PsiElement> result = import_prc.getResult();
    if (result.size() > 0) {
      for (PsiElement stmt : import_prc.getResult()) {
        for (PyImportElement ielt : ((PyImportStatement)stmt).getImportElements()) {
          final PyReferenceExpression src = ielt.getImportReference();
          if (src != null) {
            PsiElement dst = src.getReference().resolve();
            if (dst instanceof PyFile) {
              PyFile dst_file = (PyFile)dst;
              String name = ielt.getImportReference().getReferencedName(); // ref is ok or matching would fail
              seen_file_names.add(name);
              PsiElement res = dst_file.findExportedName(ref_text);
              if (res != null && !(res instanceof PyFile) && !(res instanceof PyImportElement) && dst_file.equals(res.getContainingFile())) {
                existing_import_file = dst_file;
                fix.addImport(res, dst_file, ielt);
              }
            }
          }
        }
      }
    }

    // maybe some unimported file has it, too
    ProgressManager.checkCanceled(); // before expensive index searches
    // NOTE: current indices have limitations, only finding direct definitions of classes and functions.
    Project project = node.getProject();
    List<PsiElement> symbols = new ArrayList<PsiElement>();
    symbols.addAll(PyClassNameIndex.find(ref_text, project, true));
    GlobalSearchScope scope = PyClassNameIndex.projectWithLibrariesScope(project);
    if (!isQualifier(node)) {
      symbols.addAll(PyFunctionNameIndex.find(ref_text, project, scope));
    }
    symbols.addAll(PyVariableNameIndex.find(ref_text, project, scope));
    if (!isCall(node)) {
      symbols.addAll(findImportableModules(node.getContainingFile(), ref_text, project, scope));
    }
    if (symbols.size() > 0) {
      for (PsiElement symbol : symbols) {
        if (isTopLevel(symbol)) { // we only want top-level symbols
          PsiFileSystemItem srcfile = symbol instanceof PsiFileSystemItem ? ((PsiFileSystemItem)symbol).getParent() : symbol.getContainingFile();
          if (srcfile != null && srcfile != existing_import_file && srcfile != node.getContainingFile() &&
              (ImportFromExistingAction.isRoot(project, srcfile) || PyNames.isIdentifier(FileUtil.getNameWithoutExtension(srcfile.getName())))) {
            PyQualifiedName import_path = ResolveImportUtil.findCanonicalImportPath(srcfile, node);
            if (import_path != null && !seen_file_names.contains(import_path.toString())) {
              // a new, valid hit
              fix.addImport(symbol, srcfile, import_path, proposeAsName(node.getContainingFile(), ref_text, import_path));
              seen_file_names.add(import_path.toString()); // just in case, again
            }
          }
        }
      }
    }
    for(PyImportCandidateProvider provider: Extensions.getExtensions(PyImportCandidateProvider.EP_NAME)) {
      provider.addImportCandidates(reference, ref_text, fix);
    }
    if (fix.getCandidatesCount() > 0) {
      fix.sortCandidates();
      return fix;
    }
    return null;
  }

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

  private static boolean isCall(PyElement node) {
    if (node.getParent() instanceof PyCallExpression) {
      return node == ((PyCallExpression) node.getParent()).getCallee();
    }
    return false;
  }

  private static Collection<PsiElement> findImportableModules(PsiFile targetFile, String reftext, Project project, GlobalSearchScope scope) {
    List<PsiElement> result = new ArrayList<PsiElement>();
    PsiFile[] files = FilenameIndex.getFilesByName(project, reftext + ".py", scope);
    for (PsiFile file : files) {
      PsiDirectory parent = file.getParent();
      if (parent != null && file != targetFile &&
          (parent.findFile(PyNames.INIT_DOT_PY) != null ||
           ImportFromExistingAction.isRoot(project, parent) ||
           parent == targetFile.getParent())) {
        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) {
{'fraction_non_alphanumeric': 0.10443979986066249, 'fraction_numerical': 0.001456710367977706, 'mean_word_length': 3.642752131725963, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '833266', 'n_tokens_mistral': 4455, 'n_tokens_neox': 4400, 'n_words': 975}