Improve resolve/completion for package attributes (PY-14385, PY-14387, PY-14388, PY-14335)

* Outside of import statement do not suggest names of not imported
submodules for namespace packages (PY-14385). It has been done so for
normal packages (with __init__.py) already, now the same policy
applies to namespace packages.
* While searching for imported submodules, use only immediate
children of package, not arbitrary descendants (PY-14387).
* Names defined in __init__.py for indirectly imported (intermediate)
packages are included in completion list (PY-14388).
* Removed custom collecting of submodule names from PyModuleType,
use only PyModuleType#collectImportedSubmodules for this purpose.
As result duplicate undecorated entries are no longer included
in completion list for namespace packages (PY-14335).

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -77,6 +77,7 @@ import com.jetbrains.python.refactoring.classes.PyDependenciesComparator;
 import com.jetbrains.python.refactoring.classes.extractSuperclass.PyExtractSuperclassHelper;
 import com.jetbrains.python.refactoring.classes.membersManager.PyMemberInfo;
 import com.jetbrains.python.sdk.PythonSdkType;
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -992,8 +993,9 @@ public class PyUtil {
     } // don't touch non-dirs
   }
 
+  @Contract("null -> null; !null -> !null")
   @Nullable
-  public static PsiElement turnInitIntoDir(PsiElement target) {
+  public static PsiElement turnInitIntoDir(@Nullable PsiElement target) {
     if (target instanceof PyFile && isPackage((PsiFile)target)) {
       return ((PsiFile)target).getContainingDirectory();
     }

python/src/com/jetbrains/python/psi/types/PyImportedModuleType.java

@@ -15,17 +15,20 @@
  */
 package com.jetbrains.python.psi.types;
 
+import com.google.common.collect.Sets;
 import com.intellij.codeInsight.lookup.LookupElement;
-import com.intellij.codeInsight.lookup.LookupElementBuilder;
+import com.intellij.openapi.util.Condition;
 import com.intellij.psi.PsiDirectory;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.ProcessingContext;
-import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
-import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
-import com.jetbrains.python.psi.*;
+import com.intellij.util.containers.ContainerUtil;
+import com.jetbrains.python.psi.AccessDirection;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.psi.PyUtil;
 import com.jetbrains.python.psi.impl.PyImportedModule;
+import com.jetbrains.python.psi.resolve.PointInImport;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.RatedResolveResult;
 import com.jetbrains.python.psi.resolve.ResolveImportUtil;
@@ -35,6 +38,7 @@ import org.jetbrains.annotations.Nullable;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 
 /**
  * @author yole
@@ -58,33 +62,37 @@ public class PyImportedModuleType implements PyType {
       return new PyModuleType(file, myImportedModule).resolveMember(name, location, direction, resolveContext);
     }
     else if (resolved instanceof PsiDirectory) {
-      final List<PsiElement> elements = Collections.singletonList(ResolveImportUtil.resolveChild(resolved, name, null, true, true));
+      List<PsiElement> elements = Collections.singletonList(ResolveImportUtil.resolveChild(resolved, name, null, true, true));
+      if (location != null && ResolveImportUtil.getPointInImport(location) == PointInImport.NONE) {
+        final Set<PsiElement> imported = Sets.newHashSet(PyModuleType.collectImportedSubmodules((PsiDirectory)resolved, location));
+        elements = ContainerUtil.filter(elements, new Condition<PsiElement>() {
+          @Override
+          public boolean value(PsiElement element) {
+            return imported.contains(element);
+          }
+        });
+      }
       return ResolveImportUtil.rateResults(elements);
     }
     return null;
   }
 
   public Object[] getCompletionVariants(String completionPrefix, PsiElement location, ProcessingContext context) {
-    List<LookupElement> result = new ArrayList<LookupElement>();
-    ScopeOwner scopeOwner = ScopeUtil.getScopeOwner(location);
-    assert scopeOwner != null;
-    final List<PyImportElement> importTargets = PyModuleType.getVisibleImports(scopeOwner);
-    final int imported = myImportedModule.getImportedPrefix().getComponentCount();
-    for (PyImportElement importTarget : importTargets) {
-      final QualifiedName qName = importTarget.getImportedQName();
-      if (qName != null && qName.matchesPrefix(myImportedModule.getImportedPrefix())) {
-        final List<String> components = qName.getComponents();
-        if (components.size() > imported) {
-          String module = components.get(imported);
-          result.add(LookupElementBuilder.create(module));
-        }
-      }
-    }
+    final List<LookupElement> result = new ArrayList<LookupElement>();
     final PsiElement resolved = myImportedModule.resolve();
-    if (resolved instanceof PsiDirectory) {
+    if (resolved instanceof PyFile) {
+      final PyModuleType moduleType = new PyModuleType((PyFile)resolved, myImportedModule);
+      result.addAll(moduleType.getCompletionVariantsAsLookupElements(location, context, false, false));
+    }
+    else if (resolved instanceof PsiDirectory) {
       final PsiDirectory dir = (PsiDirectory)resolved;
       if (PyUtil.isPackage(dir, location)) {
-        result.addAll(PyModuleType.getSubModuleVariants(dir, location, null));
+        if (ResolveImportUtil.getPointInImport(location) != PointInImport.NONE) {
+          result.addAll(PyModuleType.getSubModuleVariants(dir, location, null));
+        }
+        else {
+          result.addAll(PyModuleType.collectImportedSubmodulesAsLookupElements(dir, location, context.get(CTX_NAMES)));
+        }
       }
     }
     return ArrayUtil.toObjectArray(result);

python/src/com/jetbrains/python/psi/types/PyModuleType.java

@@ -29,7 +29,9 @@ import com.intellij.psi.*;
 import com.intellij.psi.scope.PsiScopeProcessor;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.QualifiedName;
+import com.intellij.util.Function;
 import com.intellij.util.ProcessingContext;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.PyCustomMember;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
@@ -331,7 +333,7 @@ public class PyModuleType implements PyType { // Modules don't descend from obje
         result.addAll(getSubModuleVariants(myModule.getContainingDirectory(), location, namesAlready));
       }
       else {
-        addImportedSubmodules(location, namesAlready, result);
+        result.addAll(collectImportedSubmodulesAsLookupElements(myModule, location, namesAlready));
       }
     }
     return result;
@@ -355,27 +357,62 @@ public class PyModuleType implements PyType { // Modules don't descend from obje
     return processor;
   }
 
-  private void addImportedSubmodules(PsiElement location, Set<String> existingNames, List<LookupElement> result) {
-    PsiFile file = location.getContainingFile();
-    if (file instanceof PyFile) {
-      PyFile pyFile = (PyFile)file;
-      PsiElement moduleBase = PyUtil.isPackage(myModule) ? myModule.getContainingDirectory() : myModule;
-      for (PyImportElement importElement : pyFile.getImportTargets()) {
-        PsiElement target = PyUtil.turnInitIntoDir(importElement.resolve());
-        if (target != null && PsiTreeUtil.isAncestor(moduleBase, target, true)) {
-          LookupElement element = null;
-          if (target instanceof PsiFileSystemItem) {
-            element = buildFileLookupElement((PsiFileSystemItem) target, existingNames);
-          }
-          else if (target instanceof PsiNamedElement) {
-            element = LookupElementBuilder.createWithIcon((PsiNamedElement)target);
-          }
-          if (element != null) {
-            result.add(element);
-          }
-        }
-      }
+  @NotNull
+  public static List<LookupElement> collectImportedSubmodulesAsLookupElements(@NotNull PsiFileSystemItem pyPackage,
+                                                                              @NotNull PsiElement location,
+                                                                              @Nullable final Set<String> existingNames) {
+
+    return ContainerUtil.mapNotNull(collectImportedSubmodules(pyPackage, location),
+                                    new Function<PsiElement, LookupElement>() {
+                                      @Override
+                                      public LookupElement fun(PsiElement element) {
+                                        if (element instanceof PsiFileSystemItem) {
+                                          return buildFileLookupElement((PsiFileSystemItem)element, existingNames);
+                                        }
+                                        else if (element instanceof PsiNamedElement) {
+                                          return LookupElementBuilder.createWithIcon((PsiNamedElement)element);
+                                        }
+                                        return null;
+                                      }
+                                    });
+  }
+
+  @NotNull
+  public static List<PsiElement> collectImportedSubmodules(@NotNull PsiFileSystemItem pyPackage, @NotNull PsiElement location) {
+    final PsiElement parentAnchor;
+    if (pyPackage instanceof PyFile && PyUtil.isPackage(((PyFile)pyPackage))) {
+      parentAnchor = ((PyFile)pyPackage).getContainingDirectory();
     }
+    else if (pyPackage instanceof PsiDirectory && PyUtil.isPackage(((PsiDirectory)pyPackage), location)) {
+      parentAnchor = pyPackage;
+    }
+    else {
+      throw new IllegalArgumentException("Package argument should be either standard Python package or namespace package");
+    }
+
+    final ScopeOwner scopeOwner = ScopeUtil.getScopeOwner(location);
+    if (scopeOwner == null) {
+      return Collections.emptyList();
+    }
+    final List<PsiElement> result = new ArrayList<PsiElement>();
+    nextImportElement:
+    for (PyImportElement importElement : getVisibleImports(scopeOwner)) {
+      PsiElement resolvedChild = PyUtil.turnInitIntoDir(importElement.resolve());
+      if (resolvedChild == null || !PsiTreeUtil.isAncestor(parentAnchor, resolvedChild, true)) {
+        continue;
+      }
+      QualifiedName importedQName = importElement.getImportedQName();
+      // Looking for strict child of parentAncestor
+      while (resolvedChild != null && resolvedChild.getParent() != parentAnchor) {
+        if (importedQName == null || importedQName.getComponentCount() <= 1) {
+          continue nextImportElement;
+        }
+        importedQName = importedQName.removeTail(1);
+        resolvedChild = PyUtil.turnInitIntoDir(ResolveImportUtil.resolveImportElement(importElement, importedQName));
+      }
+      ContainerUtil.addIfNotNull(result, resolvedChild);
+    }
+    return result;
   }
 
   public static List<LookupElement> getSubModuleVariants(final PsiDirectory directory,

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

@@ -0,0 +1,3 @@
+import pkg1.m1
+
+pkg1.foo

python/testData/completion/attributeOfIndirectlyImportedPackage/a.py

@@ -0,0 +1,3 @@
+import pkg1.m1
+
+pkg1.f<caret>

python/testData/completion/attributeOfIndirectlyImportedPackage/pkg1/__init__.py

@@ -0,0 +1 @@
+foo = 42

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

@@ -0,0 +1,3 @@
+import nspkg1.m2
+
+nspkg1.m2

python/testData/completion/notImportedSubmodulesOfNamespacePackage/a.py

@@ -0,0 +1,3 @@
+import nspkg1.m2
+
+nspkg1.m<caret>

python/testData/completion/submoduleOfIndirectlyImportedPackage/a.py

@@ -0,0 +1,4 @@
+import pkg1.subpkg1.m1
+import pkg1
+
+pkg1.<caret>

python/testData/completion/submoduleOfIndirectlyImportedPackage/pkg1/__init__.py

@@ -0,0 +1 @@
+VAR = 42

python/testData/inspections/PyUnresolvedReferencesInspection3K/NotImportedSubmodulesOfNamespacePackage/main.py

@@ -0,0 +1,5 @@
+import nspkg1.m2
+
+print(nspkg1.m2)
+print(nspkg1.<warning descr="Cannot find reference 'm3' in 'imported module nspkg1'">m3</warning>)
+print(nspkg1.<warning descr="Cannot find reference 'nssubpkg1' in 'imported module nspkg1'">nssubpkg1</warning>)

python/testSrc/com/jetbrains/python/Py3CompletionTest.java

@@ -136,4 +136,9 @@ public class Py3CompletionTest extends PyTestCase {
     myFixture.completeBasic();
     myFixture.checkResultByFile("completion/" + getTestName(true) + "/nspkg1/a.after.py");
   }
+
+  // PY-14385
+  public void testNotImportedSubmodulesOfNamespacePackage() {
+    doMultiFileTest();
+  }
 }

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

@@ -707,4 +707,19 @@ public class PythonCompletionTest extends PyTestCase {
     assertContainsElements(suggested, PyNames.FUNCTION_SPECIAL_ATTRIBUTES);
     assertDoesntContain(suggested, PyNames.METHOD_SPECIAL_ATTRIBUTES);
   }
+
+  // PY-14388
+  public void testAttributeOfIndirectlyImportedPackage() {
+    doMultiFileTest();
+  }
+
+  // PY-14387
+  public void testSubmoduleOfIndirectlyImportedPackage() {
+    myFixture.copyDirectoryToProject("completion/" + getTestName(true), "");
+    myFixture.configureByFile("a.py");
+    myFixture.completeBasic();
+    final List<String> suggested = myFixture.getLookupElementStrings();
+    assertNotNull(suggested);
+    assertSameElements(suggested, "VAR", "subpkg1");
+  }
 }

python/testSrc/com/jetbrains/python/inspections/Py3UnresolvedReferencesInspectionTest.java

@@ -111,4 +111,9 @@ public class Py3UnresolvedReferencesInspectionTest extends PyTestCase {
   public void testAttributesOfUnresolvedTypeFile() {
     doTest();
   }
+
+  // PY-14385
+  public void testNotImportedSubmodulesOfNamespacePackage() {
+    doMultiFileTest("main.py");
+  }
 }