PY-17810 PY-28989 Namespace package modules added to autocompletion

Packages are added to completion results. Common functionality extracted to PyExtendedCompletionContributor.

python/src/META-INF/python-core-common.xml

@@ -475,7 +475,7 @@
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyClassNameCompletionContributor"/>
     <completion.contributor language="Python"
-                            implementationClass="com.jetbrains.python.codeInsight.completion.PyModuleCompletionContributor"/>
+                            implementationClass="com.jetbrains.python.codeInsight.completion.PyModulePackageCompletionContributor"/>
     <weigher key="completion" implementationClass="com.jetbrains.python.codeInsight.completion.PythonCompletionWeigher" order="first"/>
     <completion.confidence language="Python" implementationClass="com.jetbrains.python.codeInsight.completion.PyCompletionConfidence"/>
     <typedHandler implementation="com.jetbrains.python.console.completion.PythonConsoleAutopopupBlockingHandler" id="pydevBlockAutoPopup"

python/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.java

@@ -6,18 +6,16 @@ import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.codeInsight.lookup.LookupElementBuilder;
 import com.intellij.navigation.ItemPresentation;
 import com.intellij.navigation.NavigationItem;
-import com.intellij.openapi.command.WriteCommandAction;
-import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Condition;
-import com.intellij.psi.*;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiNamedElement;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.stubs.StubIndex;
 import com.intellij.psi.stubs.StubIndexKey;
-import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.containers.ContainerUtil;
-import com.jetbrains.python.codeInsight.imports.AddImportHelper;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.search.PyProjectScopeBuilder;
 import com.jetbrains.python.psi.stubs.PyClassNameIndex;
@@ -32,55 +30,42 @@ import java.util.Set;
 import java.util.function.Function;
 
 /**
+ * Adds completion variants for Python classes, functions and variables.
+ *
  * @author yole
  */
-public class PyClassNameCompletionContributor extends CompletionContributor {
+public class PyClassNameCompletionContributor extends PyExtendedCompletionContributor {
 
   @Override
-  public void fillCompletionVariants(@NotNull CompletionParameters parameters, @NotNull CompletionResultSet result) {
-    if (parameters.isExtendedCompletion()) {
-      final PsiElement element = parameters.getPosition();
-      final PsiElement parent = element.getParent();
-      if (parent instanceof PyReferenceExpression && ((PyReferenceExpression)parent).isQualified()) {
-        return;
-      }
-      if (parent instanceof PyStringLiteralExpression) {
-        final String prefix = parent.getText().substring(0, parameters.getOffset() - parent.getTextRange().getStartOffset());
-        if (prefix.contains(".")) {
-          return;
-        }
-      }
-      final FileViewProvider provider = element.getContainingFile().getViewProvider();
-      if (provider instanceof MultiplePsiFilesPerDocumentFileViewProvider) return;
-      if (PsiTreeUtil.getParentOfType(element, PyImportStatementBase.class) != null) {
-        return;
-      }
-      final PsiFile originalFile = parameters.getOriginalFile();
+  protected void doFillCompletionVariants(@NotNull CompletionParameters parameters, @NotNull CompletionResultSet result) {
+    final PsiFile originalFile = parameters.getOriginalFile();
+    final PsiElement element = parameters.getPosition();
+    final PsiElement parent = element.getParent();
 
-      addVariantsFromIndex(result,
-                           originalFile,
-                           PyClassNameIndex.KEY,
-                           parent instanceof PyStringLiteralExpression ? STRING_LITERAL_INSERT_HANDLER : IMPORTING_INSERT_HANDLER,
-                           IS_TOPLEVEL,
-                           PyClass.class,
-                           createClassElementHandler(originalFile));
+    addVariantsFromIndex(result,
+                         originalFile,
+                         PyClassNameIndex.KEY,
+                         parent instanceof PyStringLiteralExpression ? getStringLiteralInsertHandler() : getImportingInsertHandler(),
+                         // TODO: implement autocompletion for inner classes
+                         PyUtil::isTopLevel,
+                         PyClass.class,
+                         createClassElementHandler(originalFile));
 
-      addVariantsFromIndex(result,
-                           originalFile,
-                           PyFunctionNameIndex.KEY,
-                           getFunctionInsertHandler(parent),
-                           IS_TOPLEVEL,
-                           PyFunction.class,
-                           Function.identity());
+    addVariantsFromIndex(result,
+                         originalFile,
+                         PyFunctionNameIndex.KEY,
+                         getFunctionInsertHandler(parent),
+                         PyUtil::isTopLevel,
+                         PyFunction.class,
+                         Function.identity());
 
-      addVariantsFromIndex(result,
-                           originalFile,
-                           PyVariableNameIndex.KEY,
-                           parent instanceof PyStringLiteralExpression ? STRING_LITERAL_INSERT_HANDLER : IMPORTING_INSERT_HANDLER,
-                           IS_TOPLEVEL,
-                           PyTargetExpression.class,
-                           Function.identity());
-    }
+    addVariantsFromIndex(result,
+                         originalFile,
+                         PyVariableNameIndex.KEY,
+                         parent instanceof PyStringLiteralExpression ? getStringLiteralInsertHandler() : getImportingInsertHandler(),
+                         PyUtil::isTopLevel,
+                         PyTargetExpression.class,
+                         Function.identity());
   }
 
   @NotNull
@@ -104,18 +89,16 @@ public class PyClassNameCompletionContributor extends CompletionContributor {
     };
   }
 
-  private static InsertHandler<LookupElement> getFunctionInsertHandler(PsiElement parent) {
+  private InsertHandler<LookupElement> getFunctionInsertHandler(PsiElement parent) {
     if (parent instanceof PyStringLiteralExpression) {
-      return STRING_LITERAL_INSERT_HANDLER;
+      return getStringLiteralInsertHandler();
     }
     if (parent.getParent() instanceof PyDecorator) {
-      return IMPORTING_INSERT_HANDLER;
+      return getImportingInsertHandler();
     }
-    return FUNCTION_INSERT_HANDLER;
+    return getFunctionInsertHandler();
   }
 
-  private static final Condition<PsiElement> IS_TOPLEVEL = element -> PyUtil.isTopLevel(element);
-
   private static <T extends PsiNamedElement> void addVariantsFromIndex(@NotNull CompletionResultSet resultSet,
                                                                        @NotNull PsiFile targetFile,
                                                                        @NotNull StubIndexKey<String, T> key,
@@ -150,51 +133,4 @@ public class PyClassNameCompletionContributor extends CompletionContributor {
                  .map(elementHandler)
                  .forEach(resultSet::addElement);
   }
-
-  static final InsertHandler<LookupElement> IMPORTING_INSERT_HANDLER = new InsertHandler<LookupElement>() {
-    public void handleInsert(final InsertionContext context, final LookupElement item) {
-        addImportForLookupElement(context, item, context.getTailOffset() - 1);
-    }
-  };
-
-
-  static final InsertHandler<LookupElement> FUNCTION_INSERT_HANDLER = new PyFunctionInsertHandler() {
-    public void handleInsert(@NotNull final InsertionContext context, @NotNull final LookupElement item) {
-      int tailOffset = context.getTailOffset()-1;
-      super.handleInsert(context, item);  // adds parentheses, modifies tail offset
-      context.commitDocument();
-      addImportForLookupElement(context, item, tailOffset);
-    }
-  };
-
-  static final InsertHandler<LookupElement> STRING_LITERAL_INSERT_HANDLER = new InsertHandler<LookupElement>() {
-    @Override
-    public void handleInsert(InsertionContext context, LookupElement item) {
-      PsiElement element = item.getPsiElement();
-      if (element instanceof PyQualifiedNameOwner) {
-        String qName = ((PyQualifiedNameOwner) element).getQualifiedName();
-        String name = ((PyQualifiedNameOwner) element).getName();
-        if (qName != null && name != null) {
-          String qNamePrefix = qName.substring(0, qName.length()-name.length());
-          context.getDocument().insertString(context.getStartOffset(), qNamePrefix);
-        }
-      }
-    }
-  };
-
-  private static void addImportForLookupElement(final InsertionContext context, final LookupElement item, final int tailOffset) {
-    PsiDocumentManager manager = PsiDocumentManager.getInstance(context.getProject());
-    Document document = manager.getDocument(context.getFile());
-    if (document != null) {
-      manager.commitDocument(document);
-    }
-    final PsiReference ref = context.getFile().findReferenceAt(tailOffset);
-    if (ref == null || ref.resolve() == item.getPsiElement()) {
-      // no import statement needed
-      return;
-    }
-    WriteCommandAction.writeCommandAction(context.getProject(), context.getFile()).run(() -> {
-      AddImportHelper.addImport(PyUtil.as(item.getPsiElement(), PsiNamedElement.class), context.getFile(), (PyElement)ref.getElement());
-    });
-  }
 }

python/src/com/jetbrains/python/codeInsight/completion/PyExtendedCompletionContributor.kt

@@ -0,0 +1,103 @@
+// 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.python.codeInsight.completion
+
+import com.intellij.codeInsight.completion.*
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.psi.MultiplePsiFilesPerDocumentFileViewProvider
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiNamedElement
+import com.intellij.psi.util.PsiTreeUtil
+import com.jetbrains.python.codeInsight.imports.AddImportHelper
+import com.jetbrains.python.psi.*
+
+/**
+ * Provides basic functionality for extended completion.
+ *
+ * Extended code completion is actually a basic code completion that shows the names of classes, functions, modules and variables.
+ *
+ * To provide variants for extended completion override [doFillCompletionVariants]
+ */
+abstract class PyExtendedCompletionContributor : CompletionContributor() {
+
+  protected val importingInsertHandler: InsertHandler<LookupElement> = InsertHandler { context, item ->
+    addImportForLookupElement(context, item, context.tailOffset - 1)
+  }
+
+  protected val functionInsertHandler: InsertHandler<LookupElement> = object : PyFunctionInsertHandler() {
+    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+      val tailOffset = context.tailOffset - 1
+      super.handleInsert(context, item)  // adds parentheses, modifies tail offset
+      context.commitDocument()
+      addImportForLookupElement(context, item, tailOffset)
+    }
+  }
+
+  protected val stringLiteralInsertHandler: InsertHandler<LookupElement> = InsertHandler { context, item ->
+    val element = item.psiElement
+    if (element is PyQualifiedNameOwner) {
+      val qName = element.qualifiedName
+      val name = element.name
+      if (qName != null && name != null) {
+        val qNamePrefix = qName.substring(0, qName.length - name.length)
+        context.document.insertString(context.startOffset, qNamePrefix)
+      }
+    }
+  }
+
+  /**
+   * Checks whether completion should be performed for a given [parameters] and delegates actual work to [doFillCompletionVariants].
+   */
+  final override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
+    if (!shouldDoCompletion(parameters)) return
+    doFillCompletionVariants(parameters, result)
+  }
+
+  /**
+   * Subclasses should override the method to provide completion variants.
+   */
+  protected abstract fun doFillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet)
+
+  private fun shouldDoCompletion(parameters: CompletionParameters): Boolean {
+    if (!parameters.isExtendedCompletion) {
+      return false
+    }
+
+    val element = parameters.position
+    val parent = element.parent
+    if (parent is PyReferenceExpression && parent.isQualified) {
+      return false
+    }
+    if (parent is PyStringLiteralExpression) {
+      val prefix = parent.text.substring(0, parameters.offset - parent.textRange.startOffset)
+      if (prefix.contains(".")) {
+        return false
+      }
+    }
+    val provider = element.containingFile.viewProvider
+    if (provider is MultiplePsiFilesPerDocumentFileViewProvider) {
+      return false
+    }
+
+    return PsiTreeUtil.getParentOfType(element, PyImportStatementBase::class.java) == null
+  }
+}
+
+private fun addImportForLookupElement(context: InsertionContext, item: LookupElement, tailOffset: Int) {
+  val manager = PsiDocumentManager.getInstance(context.project)
+  val document = manager.getDocument(context.file)
+  if (document != null) {
+    manager.commitDocument(document)
+  }
+  val ref = context.file.findReferenceAt(tailOffset)
+  if (ref == null || ref.resolve() === item.psiElement) {
+    // no import statement needed
+    return
+  }
+  WriteCommandAction.writeCommandAction(context.project, context.file).run<RuntimeException> {
+    val psiElement = item.psiElement
+    if (psiElement is PsiNamedElement) {
+      AddImportHelper.addImport(psiElement, context.file, ref.element as PyElement)
+    }
+  }
+}

python/src/com/jetbrains/python/codeInsight/completion/PyModulePackageCompletionContributor.kt

@@ -1,25 +1,28 @@
 // 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.python.codeInsight.completion
 
-import com.intellij.codeInsight.completion.CompletionContributor
 import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
-import com.jetbrains.python.codeInsight.completion.PyClassNameCompletionContributor.IMPORTING_INSERT_HANDLER
-import com.jetbrains.python.codeInsight.completion.PyClassNameCompletionContributor.STRING_LITERAL_INSERT_HANDLER
+import com.intellij.psi.PsiFileSystemItem
 import com.jetbrains.python.codeInsight.imports.PythonImportUtils
 import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.psi.PyStringLiteralExpression
-import com.jetbrains.python.psi.impl.PyPsiFacadeImpl
 import com.jetbrains.python.psi.resolve.PyQualifiedNameResolveContext
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder
+import com.jetbrains.python.psi.resolve.fromFoothold
 import com.jetbrains.python.psi.resolve.resolveQualifiedName
 import com.jetbrains.python.psi.stubs.PyModuleNameIndex
 import com.jetbrains.python.psi.types.PyModuleType
 
-class PyModuleCompletionContributor : CompletionContributor() {
+/**
+ * Add completion variants for modules and packages.
+ *
+ * The completion contributor ensures that completion variants are resolvable with project source root configuration.
+ * The list of completion variants does not include namespace packages (but includes their modules where appropriate).
+ */
+class PyModulePackageCompletionContributor : PyExtendedCompletionContributor() {
 
-  override fun fillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
-    if (!parameters.isExtendedCompletion) return
+  override fun doFillCompletionVariants(parameters: CompletionParameters, result: CompletionResultSet) {
 
     val targetFile = parameters.originalFile
     val inStringLiteral = parameters.position.parent is PyStringLiteralExpression
@@ -27,23 +30,24 @@ class PyModuleCompletionContributor : CompletionContributor() {
     val modulesFromIndex = moduleKeys.asSequence()
       .filter { result.prefixMatcher.prefixMatches(it) }
       .flatMap { PyModuleNameIndex.find(it, targetFile.project, true).asSequence() }
+      .toList()
 
-    val resolveContext = PyPsiFacadeImpl(targetFile.project).createResolveContextFromFoothold(targetFile)
-    val builders = modulesFromIndex
+    val resolveContext = fromFoothold(targetFile)
+    val builders = modulesFromIndex.asSequence()
       .flatMap { resolve(it, resolveContext) }
       .filter { PythonImportUtils.isImportableModule(targetFile, it) }
       .mapNotNull { PyModuleType.buildFileLookupElement(it, null) }
       .map { it.withInsertHandler(
-            if (inStringLiteral) STRING_LITERAL_INSERT_HANDLER else IMPORTING_INSERT_HANDLER)
+            if (inStringLiteral) stringLiteralInsertHandler else importingInsertHandler)
       }
 
     builders.forEach { result.addElement(it) }
   }
 
-  private fun resolve(module: PyFile, resolveContext: PyQualifiedNameResolveContext): Sequence<PyFile> {
+  private fun resolve(module: PyFile, resolveContext: PyQualifiedNameResolveContext): Sequence<PsiFileSystemItem> {
     val qualifiedName = QualifiedNameFinder.findCanonicalImportPath(module, null) ?: return emptySequence()
     return resolveQualifiedName(qualifiedName, resolveContext).asSequence()
-      .filterIsInstance<PyFile>()
+      .filterIsInstance<PsiFileSystemItem>()
   }
 
-}
+}

python/src/com/jetbrains/python/codeInsight/imports/PythonImportUtils.java

@@ -34,6 +34,7 @@ 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.PyResolveImportUtil;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
 import com.jetbrains.python.psi.search.PyProjectScopeBuilder;
 import com.jetbrains.python.psi.stubs.PyClassNameIndex;
@@ -254,9 +255,9 @@ public final class PythonImportUtils {
   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());
+           (ImportFromExistingAction.isRoot(parent) ||
+            parent == targetFile.getParent() ||
+            PyUtil.isPackage(parent, false, null));
   }
 
   private static boolean isIndexableTopLevel(PsiElement symbol) {

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

@@ -301,8 +301,9 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
       @Override
       public String getLocationString() {
         final PyClass containingClass = getContainingClass();
-        if (containingClass != null) {
-          return "(" + containingClass.getName() + " in " + getPackageForFile(getContainingFile()) + ")";
+        final String packageForFile = getPackageForFile(getContainingFile());
+        if (containingClass != null && packageForFile != null) {
+          return String.format("(%s in %s)", containingClass.getName(), packageForFile);
         }
         return super.getLocationString();
       }

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

@@ -341,7 +341,8 @@ public class PyStringLiteralExpressionImpl extends PyElementImpl implements PySt
       @Nullable
       @Override
       public String getLocationString() {
-        return "(" + PyElementPresentation.getPackageForFile(getContainingFile()) + ")";
+        String packageForFile = PyElementPresentation.getPackageForFile(getContainingFile());
+        return packageForFile != null ? String.format("(%s)", packageForFile) : null;
       }
 
       @Nullable

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

@@ -705,8 +705,9 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
       @Override
       public String getLocationString() {
         final PyClass containingClass = getContainingClass();
-        if (containingClass != null) {
-          return "(" + containingClass.getName() + " in " + getPackageForFile(getContainingFile()) + ")";
+        final String packageForFile = getPackageForFile(getContainingFile());
+        if (containingClass != null && packageForFile != null) {
+          return String.format("(%s in %s)", containingClass.getName(), packageForFile);
         }
         return super.getLocationString();
       }

python/testData/completion/duplicatedEntriesFromMultipleSourceRoots/a.after.py

@@ -0,0 +1,3 @@
+from m1 import Cls
+
+Cls

python/testData/completion/fromPackageImport/a.py

@@ -0,0 +1 @@
+from pkg import <caret>

python/testData/completion/moduleFromMultipleSourceRoots/a.after.py

@@ -0,0 +1,3 @@
+import m7
+
+m7

python/testData/completion/moduleFromNamespacePackage/a.after.py

@@ -0,0 +1,3 @@
+from pkg import m7
+
+m7

python/testData/completion/moduleFromNamespacePackage/a.py

@@ -0,0 +1 @@
+m7<caret>

python/testData/completion/packageFromMultipleSourceRoots/a.after.py

@@ -0,0 +1,3 @@
+import pkg1
+
+pkg1

python/testData/completion/packageFromMultipleSourceRoots/a.py

@@ -0,0 +1 @@
+pkg1<caret>

python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java

@@ -2,7 +2,6 @@
 package com.jetbrains.python;
 
 import com.google.common.collect.Lists;
-import com.intellij.openapi.module.Module;
 import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.roots.ModuleRootManager;
 import com.intellij.openapi.vfs.VfsUtilCore;
@@ -10,7 +9,6 @@ import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.testFramework.PlatformTestUtil;
-import com.intellij.testFramework.PsiTestUtil;
 import com.jetbrains.python.codeInsight.PyCustomMember;
 import com.jetbrains.python.fixtures.PyMultiFileResolveTestCase;
 import com.jetbrains.python.fixtures.PyResolveTestCase;
@@ -438,7 +436,7 @@ public class PyMultiFileResolveTest extends PyMultiFileResolveTestCase {
   // PY-18626
   public void testManySourceRoots() {
     myFixture.copyDirectoryToProject("manySourceRoots", "");
-    withSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {
+    runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {
       final PsiFile psiFile = myFixture.configureByFile("a.py");
       final PsiReference ref = PyResolveTestCase.findReferenceByMarker(psiFile);
       assertInstanceOf(ref, PsiPolyVariantReference.class);
@@ -454,7 +452,7 @@ public class PyMultiFileResolveTest extends PyMultiFileResolveTestCase {
   // PY-28321
   public void testImportManySourceRoots() {
     myFixture.copyDirectoryToProject("importManySourceRoots", "");
-    withSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root2"), myFixture.findFileInTempDir("root1")), () -> {
+    runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root2"), myFixture.findFileInTempDir("root1")), () -> {
       final PsiFile psiFile = myFixture.configureByFile("root1/pkg/a.py");
       final PsiReference ref = PyResolveTestCase.findReferenceByMarker(psiFile);
       assertInstanceOf(ref, PsiPolyVariantReference.class);
@@ -470,7 +468,7 @@ public class PyMultiFileResolveTest extends PyMultiFileResolveTestCase {
   // PY-28321
   public void testImportManySourceRootsReverseRootOrder() {
     myFixture.copyDirectoryToProject("importManySourceRoots", "");
-    withSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {
+    runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {
       final PsiFile psiFile = myFixture.configureByFile("root1/pkg/a.py");
       final PsiReference ref = PyResolveTestCase.findReferenceByMarker(psiFile);
       assertInstanceOf(ref, PsiPolyVariantReference.class);
@@ -485,7 +483,7 @@ public class PyMultiFileResolveTest extends PyMultiFileResolveTestCase {
 
     VirtualFile vf = myFixture.findFileInTempDir("ext/m1.py");
 
-    withSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root")), () -> {
+    runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root")), () -> {
       final PsiFile extSource = myFixture.getPsiManager().findFile(vf);
       PyImportResolver foreignResolver = (name, context, withRoots) -> name.toString().equals("m1") ? extSource : null;
       PlatformTestUtil.registerExtension(PyImportResolver.EP_NAME, foreignResolver, getTestRootDisposable());

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -10,7 +10,7 @@ import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.codeInsight.lookup.LookupElementBuilder;
 import com.intellij.testFramework.TestDataPath;
 import com.jetbrains.python.documentation.docstrings.DocStringFormat;
-import com.jetbrains.python.fixtures.PyMultiFileResolveTestCase;
+import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -21,7 +21,7 @@ import java.util.List;
 import java.util.Set;
 
 @TestDataPath("$CONTENT_ROOT/../testData/completion")
-public class PythonCompletionTest extends PyMultiFileResolveTestCase {
+public class PythonCompletionTest extends PyTestCase {
 
   private void doTest() {
     CamelHumpMatcher.forceStartMatching(myFixture.getTestRootDisposable());
@@ -32,9 +32,13 @@ public class PythonCompletionTest extends PyMultiFileResolveTestCase {
   }
 
   private void doMultiFileTest() {
+    doMultiFileTest(CompletionType.BASIC, 1);
+  }
+
+  private void doMultiFileTest(CompletionType completionType, int invocationCount) {
     myFixture.copyDirectoryToProject(getTestName(true), "");
     myFixture.configureByFile("a.py");
-    myFixture.completeBasic();
+    myFixture.complete(completionType, invocationCount);
     myFixture.checkResultByFile(getTestName(true) + "/a.after.py");
   }
 
@@ -1288,7 +1292,7 @@ public class PythonCompletionTest extends PyMultiFileResolveTestCase {
 
   // PY-17810
   public void testDuplicatedEntriesFromMultipleSourceRoots() {
-    assertSingleVariantInExtendedCompletion();
+    assertSingleVariantInExtendedCompletionWithSourceRoots();
   }
 
   // PY-17810
@@ -1298,7 +1302,27 @@ public class PythonCompletionTest extends PyMultiFileResolveTestCase {
 
   // PY-17810
   public void testModuleFromMultipleSourceRoots() {
-    assertSingleVariantInExtendedCompletion();
+    assertSingleVariantInExtendedCompletionWithSourceRoots();
+  }
+
+  // PY-17810
+  public void testPackageFromMultipleSourceRoots() {
+    assertSingleVariantInExtendedCompletionWithSourceRoots();
+  }
+
+  // PY-17810
+  public void testFromPackageImport() {
+    myFixture.copyDirectoryToProject(getTestName(true), "");
+    myFixture.configureByFile("a.py");
+    myFixture.complete(CompletionType.BASIC, 2);
+    final List<String> suggested = myFixture.getLookupElementStrings();
+    assertNotNull(suggested);
+    assertSameElements(suggested, "m1", "m2");
+  }
+
+  // PY-28989
+  public void testModuleFromNamespacePackage() {
+    runWithLanguageLevel(LanguageLevel.PYTHON34, this::assertSingleVariantInExtendedCompletion);
   }
 
   private void assertNoVariantsInExtendedCompletion() {
@@ -1310,15 +1334,20 @@ public class PythonCompletionTest extends PyMultiFileResolveTestCase {
   }
 
   private void assertSingleVariantInExtendedCompletion() {
+    doMultiFileTest(CompletionType.BASIC, 2);
+    assertNull(myFixture.getLookupElements());
+  }
+
+  private void assertSingleVariantInExtendedCompletionWithSourceRoots() {
     myFixture.copyDirectoryToProject(getTestName(true), "");
-    withSourceRoots(Lists.newArrayList(
+    runWithSourceRoots(Lists.newArrayList(
       myFixture.findFileInTempDir("root1"),
       myFixture.findFileInTempDir("root2")),
-                    () -> {
-                      myFixture.configureByFile("a.py");
-                      myFixture.complete(CompletionType.BASIC, 2);
-                      assertNull(myFixture.getLookupElements());
-                    });
+                       () -> {
+                         myFixture.configureByFile("a.py");
+                         assertNull(myFixture.complete(CompletionType.BASIC, 2));
+                         myFixture.checkResultByFile(getTestName(true) + "/a.after.py");
+                       });
   }
 
   @Override

python/testSrc/com/jetbrains/python/fixtures/PyMultiFileResolveTestCase.java

@@ -16,12 +16,10 @@
 package com.jetbrains.python.fixtures;
 
 import com.intellij.openapi.fileTypes.FileType;
-import com.intellij.openapi.module.Module;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.openapi.vfs.VirtualFileFilter;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.PsiManagerEx;
-import com.intellij.testFramework.PsiTestUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.PythonTestUtil;
@@ -100,13 +98,4 @@ public abstract class PyMultiFileResolveTestCase extends PyResolveTestCase {
     return Collections.singletonList(ref.resolve());
   }
 
-  protected void withSourceRoots(@NotNull List<VirtualFile> sourceRoots, @NotNull Runnable f) {
-    final Module module = myFixture.getModule();
-    sourceRoots.forEach(root -> PsiTestUtil.addSourceRoot(module, root));
-    try {
-      f.run();
-    } finally {
-      sourceRoots.forEach(root -> PsiTestUtil.removeSourceRoot(module, root));
-    }
-  }
 }

python/testSrc/com/jetbrains/python/fixtures/PyTestCase.java

@@ -205,10 +205,10 @@ public abstract class PyTestCase extends UsefulTestCase {
     PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), languageLevel);
   }
 
-  protected void runWithLanguageLevel(@NotNull LanguageLevel languageLevel, @NotNull Runnable action) {
+  protected void runWithLanguageLevel(@NotNull LanguageLevel languageLevel, @NotNull Runnable runnable) {
     setLanguageLevel(languageLevel);
     try {
-      action.run();
+      runnable.run();
     }
     finally {
       setLanguageLevel(null);
@@ -227,6 +227,16 @@ public abstract class PyTestCase extends UsefulTestCase {
     }
   }
 
+  protected void runWithSourceRoots(@NotNull List<VirtualFile> sourceRoots, @NotNull Runnable runnable) {
+    final Module module = myFixture.getModule();
+    sourceRoots.forEach(root -> PsiTestUtil.addSourceRoot(module, root));
+    try {
+      runnable.run();
+    } finally {
+      sourceRoots.forEach(root -> PsiTestUtil.removeSourceRoot(module, root));
+    }
+  }
+
   protected static void assertNotParsed(PsiFile file) {
     assertInstanceOf(file, PyFileImpl.class);
     assertNull("Operations should have been performed on stubs but caused file to be parsed: " + file.getVirtualFile().getPath(),