PY-28848: Complete python modules in PySymbolFieldWithBrowseButton

If we have no qualified name yet we still provide list
of all possible modules (files and packages with __init__.py) found in module or project and sdk roots.

There is also refactoring that adds extension methods both to VirtualFile and PsiFileSystemItem to check if one is python module. Former does not use latter because of performance.

This functionality is enabled for flask, tests, and plain python script

python/src/com/jetbrains/PySymbolFieldWithBrowseButton.kt

@@ -20,6 +20,7 @@ import com.intellij.util.ProcessingContext
 import com.intellij.util.TextFieldCompletionProvider
 import com.intellij.util.textCompletion.TextFieldWithCompletion
 import com.jetbrains.extensions.getQName
+import com.jetbrains.extensions.python.toPythonPsi
 import com.jetbrains.extenstions.ContextAnchor
 import com.jetbrains.extenstions.QNameResolveContext
 import com.jetbrains.extenstions.resolveToElement
@@ -38,7 +39,7 @@ import com.jetbrains.python.psi.types.TypeEvalContext
 
 /**
  * Text field to enter python symbols and browse button (from [PyGotoSymbolContributor]).
- * Supports auto-completion for symbol fully qualified names inside of textbox (except first symbol).
+ * Supports auto-completion for symbol fully qualified names inside of textbox
  * @param filter lambda to filter symbols
  * @param startFromDirectory symbols resolved against module, but may additionally be resolved against this folder if provided like in [QNameResolveContext.folderToStart]
  *
@@ -76,39 +77,53 @@ private class PyNameCompletionProvider(private val contextAnchor: ContextAnchor,
                                        private val startFromDirectory: (() -> VirtualFile)? = null) : TextFieldCompletionProvider() {
   override fun addCompletionVariants(text: String, offset: Int, prefix: String, result: CompletionResultSet) {
 
-    val evalContext = TypeEvalContext.userInitiated(contextAnchor.project, null)
-    var name = QualifiedName.fromDottedString(text)
-    val resolveContext = QNameResolveContext(contextAnchor, evalContext = evalContext, allowInaccurateResult = false,
-                                             folderToStart = startFromDirectory?.invoke())
-    var element = name.resolveToElement(resolveContext, stopOnFirstFail = true)
-
-    if (element == null && name.componentCount > 1) {
-      name = name.removeLastComponent()
-      element = name.resolveToElement(resolveContext, stopOnFirstFail = true)
-    }
-    if (element == null) {
-      return
+    val lookups: Array<LookupElement>
+    var name: QualifiedName? = null
+    if ('.' !in text) {
+      lookups = contextAnchor.getRoots()
+        .map { rootFolder -> rootFolder.children.map { it.toPythonPsi(contextAnchor.project) } }
+        .flatten()
+        .filterNotNull()
+        .toSet() // Unique
+        .map {LookupElementBuilder.create(it, it.virtualFile.nameWithoutExtension)}
+        .toTypedArray()
     }
+    else {
 
-    val lookups: Array<LookupElement> = when (element) {
-      is PyFile -> PyModuleType(element).getVariants(element)
-      is PsiDirectory -> {
-        val init = PyUtil.turnDirIntoInit(element) as? PyFile ?: return
-        PyModuleType(init).getVariants(element) +
-        element.children.filterIsInstance(PsiFileSystemItem::class.java)
-          // For package we need all symbols in initpy and all filesystem children of this folder except initpy itself
-          .filterNot { it.name == PyNames.INIT_DOT_PY }
-          .map { LookupElementBuilder.create(it, it.virtualFile.nameWithoutExtension) }
+      val evalContext = TypeEvalContext.userInitiated(contextAnchor.project, null)
+      name = QualifiedName.fromDottedString(text)
+      val resolveContext = QNameResolveContext(contextAnchor, evalContext = evalContext, allowInaccurateResult = false,
+                                               folderToStart = startFromDirectory?.invoke())
+      var element = name.resolveToElement(resolveContext, stopOnFirstFail = true)
+
+      if (element == null && name.componentCount > 1) {
+        name = name.removeLastComponent()
+        element = name.resolveToElement(resolveContext, stopOnFirstFail = true)
       }
-      is PyTypedElement -> {
-        evalContext.getType(element)?.getVariants(element) ?: return
+      if (element == null) {
+        return
+      }
+
+      lookups = when (element) {
+        is PyFile -> PyModuleType(element).getVariants(element)
+        is PsiDirectory -> {
+          val init = PyUtil.turnDirIntoInit(element) as? PyFile ?: return
+          PyModuleType(init).getVariants(element) +
+          element.children.filterIsInstance(PsiFileSystemItem::class.java)
+            // For package we need all symbols in initpy and all filesystem children of this folder except initpy itself
+            .filterNot { it.name == PyNames.INIT_DOT_PY }
+            .map { LookupElementBuilder.create(it, it.virtualFile.nameWithoutExtension) }
+        }
+        is PyTypedElement -> {
+          evalContext.getType(element)?.getVariants(element) ?: return
+        }
+        else -> return
       }
-      else -> return
     }
     result.addAllElements(lookups
                             .filter { it.psiElement != null }
                             .filter { filter?.invoke(it.psiElement!!) ?: true }
-                            .map { LookupElementBuilder.create("$name.${it.lookupString}") })
+                            .map { if (name != null) LookupElementBuilder.create("$name.${it.lookupString}") else it })
   }
 }
 

python/src/com/jetbrains/extensions/PsiFileSystemItemExt.kt

@@ -15,8 +15,11 @@
  */
 package com.jetbrains.extensions
 
+import com.intellij.psi.PsiDirectory
 import com.intellij.psi.PsiFileSystemItem
 import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.psi.PyPsiFacade
 
 /**
@@ -26,4 +29,9 @@ import com.jetbrains.python.psi.PyPsiFacade
 fun PsiFileSystemItem.getQName():QualifiedName?  {
   val name = PyPsiFacade.getInstance(this.project).findShortestImportableName(this.virtualFile, this) ?: return null
   return QualifiedName.fromDottedString(name)
-}
+}
+
+/**
+ * @return pyfile or package
+ */
+fun PsiFileSystemItem.isPythonModule() = this is PyFile || (this is PsiDirectory && this.findFile(PyNames.INIT_DOT_PY) != null)

python/src/com/jetbrains/extensions/python/VirtualFileExt.kt

@@ -0,0 +1,26 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.extensions.python
+
+import com.intellij.openapi.fileTypes.FileTypeRegistry
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiFileSystemItem
+import com.intellij.psi.PsiManager
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.PythonFileType
+
+fun VirtualFile.isPythonFile() = PythonFileType.INSTANCE == FileTypeRegistry.getInstance().getFileTypeByFileName(name)
+fun VirtualFile.isPythonPackage() = isDirectory && findChild(PyNames.INIT_DOT_PY) != null
+/**
+ * @return pyfile or package
+ */
+fun VirtualFile.isPythonModule() = isPythonPackage() || isPythonFile()
+/**
+ * @return psi element if and only if vfs item is python module
+ */
+fun VirtualFile.toPythonPsi(project: Project): PsiFileSystemItem? {
+  if (!isPythonModule()) return null
+  val manager = PsiManager.getInstance(project)
+  return if (this.isDirectory) manager.findDirectory(this) else manager.findFile(this)
+
+}

python/src/com/jetbrains/extenstions/QualifiedNameExt.kt

@@ -18,6 +18,9 @@ package com.jetbrains.extenstions
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.roots.ModuleRootManager
+import com.intellij.openapi.roots.OrderRootType
+import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiDirectory
 import com.intellij.psi.PsiElement
@@ -37,18 +40,29 @@ interface ContextAnchor {
   val project: Project
   val qualifiedNameResolveContext: PyQualifiedNameResolveContext?
   val scope: GlobalSearchScope
+  fun getRoots(): Array<VirtualFile> {
+    return sdk?.rootProvider?.getFiles(OrderRootType.CLASSES) ?: emptyArray()
+  }
 }
 
-class ModuleBasedContextAnchor(module: Module) : ContextAnchor {
+class ModuleBasedContextAnchor(val module: Module) : ContextAnchor {
   override val sdk = module.getSdk()
   override val project = module.project
   override val qualifiedNameResolveContext = fromModule(module)
   override val scope = module.moduleContentScope
+  override fun getRoots(): Array<VirtualFile> {
+    val manager = ModuleRootManager.getInstance(module)
+    return super.getRoots() + manager.contentRoots + manager.sourceRoots
+  }
 }
 
 class ProjectSdkContextAnchor(override val project: Project, override val sdk: Sdk?) : ContextAnchor {
   override val qualifiedNameResolveContext = sdk?.let { fromSdk(project, it) }
   override val scope = GlobalSearchScope.projectScope(project) //TODO: Check if project scope includes SDK
+  override fun getRoots(): Array<VirtualFile> {
+    val manager = ProjectRootManager.getInstance(project)
+    return super.getRoots() + manager.contentRoots + manager.contentSourceRoots
+  }
 }
 
 

python/src/com/jetbrains/python/psi/impl/PythonLanguageLevelPusher.java

@@ -48,6 +48,7 @@ import com.intellij.util.FileContentUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.io.DataInputOutputUtil;
 import com.intellij.util.messages.MessageBus;
+import com.jetbrains.extensions.python.VirtualFileExtKt;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.PythonModuleTypeBase;
 import com.jetbrains.python.codeInsight.typing.PyTypeShed;
@@ -209,19 +210,15 @@ public class PythonLanguageLevelPusher implements FilePropertyPusher<LanguageLev
     }
 
     if (!areLanguageLevelsCompatible(oldLanguageLevel, level) || !ProjectFileIndex.getInstance(project).isInContent(fileOrDir)) {
-      PushedFilePropertiesUpdater.getInstance(project).filePropertiesChanged(fileOrDir, PythonLanguageLevelPusher::isPythonFile);
+      PushedFilePropertiesUpdater.getInstance(project).filePropertiesChanged(fileOrDir, VirtualFileExtKt::isPythonFile);
     }
     for (VirtualFile child : fileOrDir.getChildren()) {
-      if (!child.isDirectory() && isPythonFile(child)) {
+      if (!child.isDirectory() && VirtualFileExtKt.isPythonFile(child)) {
         clearSdkPathCache(child);
       }
     }
   }
 
-  private static boolean isPythonFile(VirtualFile child) {
-    return PythonFileType.INSTANCE.equals(FileTypeRegistry.getInstance().getFileTypeByFileName(child.getName()));
-  }
-
   private static void clearSdkPathCache(@NotNull final VirtualFile child) {
     final Project[] projects = ProjectManager.getInstance().getOpenProjects();
     for (Project project : projects) {

python/src/com/jetbrains/python/run/PythonRunConfigurationForm.java

@@ -30,6 +30,7 @@ import com.intellij.openapi.util.io.FileUtil;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiFileSystemItem;
 import com.intellij.ui.PanelWithAnchor;
 import com.intellij.ui.RawCommandLineEditor;
 import com.intellij.ui.UserActivityProviderComponent;
@@ -37,12 +38,12 @@ import com.intellij.ui.components.JBCheckBox;
 import com.intellij.ui.components.JBComboBoxLabel;
 import com.intellij.ui.components.JBLabel;
 import com.jetbrains.PySymbolFieldWithBrowseButton;
+import com.jetbrains.extensions.PsiFileSystemItemExtKt;
 import com.jetbrains.extensions.python.FileChooserDescriptorExtKt;
 import com.jetbrains.extenstions.ContextAnchor;
 import com.jetbrains.extenstions.ModuleBasedContextAnchor;
 import com.jetbrains.extenstions.ProjectSdkContextAnchor;
 import com.jetbrains.python.debugger.PyDebuggerOptionsProvider;
-import com.jetbrains.python.psi.PyFile;
 import org.jetbrains.annotations.NotNull;
 
 import javax.swing.*;
@@ -117,7 +118,8 @@ public class PythonRunConfigurationForm implements PythonRunConfigurationParams,
     final ContextAnchor contentAnchor =
       (module != null ? new ModuleBasedContextAnchor(module) : new ProjectSdkContextAnchor(myProject, sdk));
     myModuleField = new PySymbolFieldWithBrowseButton(contentAnchor,
-                                                      element -> element instanceof PyFile, () -> {
+                                                      element -> element instanceof PsiFileSystemItem
+                                                                 && PsiFileSystemItemExtKt.isPythonModule((PsiFileSystemItem)element), () -> {
       final String workingDirectory = myCommonOptionsForm.getWorkingDirectory();
       if (StringUtil.isEmpty(workingDirectory)) {
         return null;