From 49dcf80e4309e2dcc3c06f2cd2e6a218a1607dde Mon Sep 17 00:00:00 2001
From: Kiryl Chetyrbak <kirylchetyrbak@gmail.com>
Date: Wed, 12 Apr 2017 15:57:17 -0400
Subject: [PATCH] Prevent PyUnresolvedReferenceInspection from scanning all
 __init__ files when generating quickfix

fixes according to review
---
 .../imports/PythonImportUtils.java            | 35 ++++++++++++-------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PythonImportUtils.java b/python/src/com/jetbrains/python/codeInsight/imports/PythonImportUtils.java
index cdfd818ca0a9..884c4118cec9 100644
--- a/python/src/com/jetbrains/python/codeInsight/imports/PythonImportUtils.java
+++ b/python/src/com/jetbrains/python/codeInsight/imports/PythonImportUtils.java
@@ -23,11 +23,14 @@ 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.openapi.vfs.VirtualFile;
 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.intellij.util.Processor;
+import com.intellij.util.indexing.FileBasedIndex;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
@@ -234,20 +237,28 @@ public final class PythonImportUtils {
 
   private static Collection<PsiElement> findImportableModules(PsiFile targetFile, String reftext, Project project, GlobalSearchScope scope) {
     List<PsiElement> result = new ArrayList<>();
-    PsiFile[] files = FilenameIndex.getFilesByName(project, reftext + ".py", scope);
-    for (PsiFile file : files) {
-      if (isImportableModule(targetFile, file)) {
-        result.add(file);
+    //add packages
+    FilenameIndex.processFilesByName(reftext, true, file -> {
+      ProgressManager.checkCanceled();
+      if (!file.isDirectory() ){
+        return true;
       }
-    }
-    // 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);
+      PsiDirectory candidatePackageDir = (PsiDirectory) file;
+      if (candidatePackageDir.findFile(PyNames.INIT_DOT_PY) == null){
+        return true;
       }
-    }
+      result.add(candidatePackageDir);
+      return true;
+    }, scope, project, null);
+    //Add modules
+    FilenameIndex.processFilesByName(reftext + ".py", false, true, item -> {
+      ProgressManager.checkCanceled();
+      if (isImportableModule(targetFile, item)){
+        result.add(item);
+      }
+      return true;
+    }, scope, project, null);
+
     return result;
   }