IDEA-CR-51114: multiresolve import references instead of single resolve (PY-32268)

There can be several resolve results for import reference. One of the examples is namespace package, when we have to consider elements from several locations for completion.

GitOrigin-RevId: 4ac824ca1f7031b2adda119684d837c6d21b3212

python/src/com/jetbrains/python/psi/impl/references/PyImportReference.java

@@ -182,21 +182,25 @@ public class PyImportReference extends PyReferenceImpl {
       if (fromImport != null && myElement.getParent() != fromImport) { // in "from foo import _"
         PyReferenceExpression src = fromImport.getImportSource();
         if (src != null) {
-          PsiElement modCandidate = src.getReference().resolve();
-          if (modCandidate instanceof PyExpression) {
-            addImportedNames(fromImport.getImportElements()); // don't propose already imported items
-            // try to collect submodules
-            PyExpression module = (PyExpression)modCandidate;
-            PyType qualifierType = myContext.getType(module);
-            if (qualifierType != null) {
-              ProcessingContext ctx = new ProcessingContext();
-              ctx.put(PyType.CTX_NAMES, myNamesAlready);
-              Collections.addAll(myObjects, qualifierType.getCompletionVariants(myElement.getName(), myElement, ctx));
+          ResolveResult[] resolved = src.getReference().multiResolve(false);
+          for (ResolveResult result : resolved) {
+            PsiElement modCandidate = result.getElement();
+            if (modCandidate instanceof PyExpression) {
+              addImportedNames(fromImport.getImportElements()); // don't propose already imported items
+              // try to collect submodules
+              PyExpression module = (PyExpression)modCandidate;
+              PyType qualifierType = myContext.getType(module);
+              if (qualifierType != null) {
+                ProcessingContext ctx = new ProcessingContext();
+                ctx.put(PyType.CTX_NAMES, myNamesAlready);
+                Collections.addAll(myObjects, qualifierType.getCompletionVariants(myElement.getName(), myElement, ctx));
+              }
+            }
+            else if (modCandidate instanceof PsiDirectory) {
+              fillFromDir((PsiDirectory)modCandidate, ImportKeywordHandler.INSTANCE);
             }
-            return myObjects.toArray();
           }
-          else if (modCandidate instanceof PsiDirectory) {
-            fillFromDir((PsiDirectory)modCandidate, ImportKeywordHandler.INSTANCE);
+          if (!myObjects.isEmpty()) {
             return myObjects.toArray();
           }
         }

python/testData/completion/fromLocalNamespacePackageInFromImport/root1/pkg/test.py

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

python/testData/completion/fromLocalNamespacePackageInImportStatement/root1/pkg/test.py

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

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

@@ -14,6 +14,7 @@ import com.intellij.testFramework.TestDataPath;
 import com.jetbrains.python.documentation.docstrings.DocStringFormat;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -619,6 +620,16 @@ public class PythonCompletionTest extends PyTestCase {
     doMultiFileTest();
   }
 
+  // PY-32268
+  public void testFromLocalNamespacePackageInFromImport() {
+    doTestImportFromLocalPython3NamespacePackage();
+  }
+
+  // PY-32268
+  public void testFromLocalNamespacePackageInImportStatement() {
+    doTestImportFromLocalPython3NamespacePackage();
+  }
+
   // PY-6829
   public void testFakeNameInQualifiedReference() {
     doTest();
@@ -1555,6 +1566,19 @@ public class PythonCompletionTest extends PyTestCase {
                        });
   }
 
+  private void doTestImportFromLocalPython3NamespacePackage() {
+    runWithLanguageLevel(LanguageLevel.PYTHON35, () -> {
+      myFixture.copyDirectoryToProject(getTestName(true), "");
+      runWithSourceRoots(Lists.newArrayList(myFixture.findFileInTempDir("root1"), myFixture.findFileInTempDir("root2")), () -> {
+        myFixture.configureByFile("root1/pkg/test.py");
+        List<String> lookupStrings = StreamEx.of(myFixture.completeBasic())
+          .map(LookupElement::getLookupString)
+          .toList();
+        assertContainsElements(lookupStrings, "foo", "bar", "baz");
+      });
+    });
+  }
+
   @Override
   protected String getTestDataPath() {
     return PythonTestUtil.getTestDataPath() + "/completion";