PY-46056 Fix wrong same directory and library elements prioritization in completion

(cherry picked from commit 0189cf588066af3c0f133caa599d1f1bc50caec0)

IJ-MR-4895

GitOrigin-RevId: d8918a1530c612138c80850e41ad140d61300f4b

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

@@ -215,6 +215,7 @@ public class PyImportReference extends PyReferenceImpl {
             }
           }
         }
+        fillFromRootsIfNotRelative(relativeLevel, insertHandler);
       }
       else { // in "import _" or "from _ import"
         PsiElement prevElem = PyPsiUtils.getPrevNonWhitespaceSibling(myElement);
@@ -235,28 +236,43 @@ public class PyImportReference extends PyReferenceImpl {
             addImportedNames(importStatement.getImportElements());
           }
         }
+        final PsiDirectory containingDirectory = myCurrentFile.getContainingDirectory();
         // look at dir by level
-        if ((relativeLevel >= 0 || !ResolveImportUtil.isAbsoluteImportEnabledFor(myCurrentFile))) {
-          final PsiDirectory containingDirectory = myCurrentFile.getContainingDirectory();
-          if (containingDirectory != null) {
-            QualifiedName thisQName = QualifiedNameFinder.findShortestImportableQName(containingDirectory);
-            if (thisQName == null || thisQName.getComponentCount() == relativeLevel) {
-              fillFromDir(ResolveImportUtil.stepBackFrom(myCurrentFile, relativeLevel), insertHandler);
-            }
-            else if (thisQName.getComponentCount() > relativeLevel) {
-              thisQName = thisQName.removeTail(relativeLevel);
-              fillFromQName(thisQName, insertHandler);
-            }
-          }
+        if (LanguageLevel.forElement(myCurrentFile).isPy3K() && containingDirectory != null &&
+            PyUtil.isExplicitPackage(containingDirectory)) {
+          fillFromRootsIfNotRelative(relativeLevel, insertHandler);
+          fillFromSameDirectoryOrRelative(relativeLevel, insertHandler);
+        }
+        else {
+          fillFromSameDirectoryOrRelative(relativeLevel, insertHandler);
+          fillFromRootsIfNotRelative(relativeLevel, insertHandler);
         }
-      }
-      if (relativeLevel == -1) {
-        fillFromQName(QualifiedName.fromComponents(), insertHandler);
       }
 
       return ArrayUtil.toObjectArray(myObjects);
     }
 
+    private void fillFromRootsIfNotRelative(int relativeLevel, @Nullable InsertHandler<LookupElement> insertHandler) {
+      if (relativeLevel == -1) {
+        fillFromQName(QualifiedName.fromComponents(), insertHandler);
+      }
+    }
+
+    private void fillFromSameDirectoryOrRelative(int relativeLevel, @Nullable InsertHandler<LookupElement> insertHandler) {
+      if (relativeLevel < 0 && ResolveImportUtil.isAbsoluteImportEnabledFor(myCurrentFile)) return;
+      final PsiDirectory containingDirectory = myCurrentFile.getContainingDirectory();
+      if (containingDirectory != null) {
+        QualifiedName thisQName = QualifiedNameFinder.findShortestImportableQName(containingDirectory);
+        if (thisQName == null || thisQName.getComponentCount() == relativeLevel) {
+          fillFromDir(ResolveImportUtil.stepBackFrom(myCurrentFile, relativeLevel), insertHandler);
+        }
+        else if (thisQName.getComponentCount() > relativeLevel) {
+          thisQName = thisQName.removeTail(relativeLevel);
+          fillFromQName(thisQName, insertHandler);
+        }
+      }
+    }
+
     private void fillFromQName(QualifiedName thisQName, InsertHandler<LookupElement> insertHandler) {
       StreamEx.of(PyResolveImportUtil.resolveQualifiedName(thisQName, PyResolveImportUtil.fromFoothold(myCurrentFile)))
         .select(PsiDirectory.class)

python/testData/completion/fromImportCompletionHintForSameDirectoryModuleInOrdinaryPackage/ordinaryPackage/sample.py

@@ -0,0 +1 @@
+from log<caret>

python/testData/completion/fromImportCompletionHintForSameDirectoryModuleInPlainDirectory/plainDirectory/sample.py

@@ -0,0 +1 @@
+from log<caret>

python/testData/completion/importCompletionHintForSameDirectoryModuleInOrdinaryPackage/ordinaryPackage/sample.py

@@ -0,0 +1 @@
+import log<caret>

python/testData/completion/importCompletionHintForSameDirectoryModuleInPlainDirectory/plainDirectory/sample.py

@@ -0,0 +1 @@
+import log<caret>

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

@@ -7,6 +7,7 @@ import com.intellij.openapi.module.Module;
 import com.intellij.testFramework.LightProjectDescriptor;
 import com.intellij.testFramework.PsiTestUtil;
 import com.intellij.testFramework.TestDataPath;
+import com.intellij.testFramework.fixtures.TestLookupElementPresentation;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.codeInsight.completion.PyModuleNameCompletionContributor;
 import com.jetbrains.python.fixtures.PyTestCase;
@@ -491,6 +492,38 @@ public class Py3CompletionTest extends PyTestCase {
     });
   }
 
+  // PY-46056
+  public void testImportCompletionHintForSameDirectoryModuleInOrdinaryPackage() {
+    doTestVariantTailText("ordinaryPackage/sample.py", "logging", null);
+  }
+
+  // PY-46056
+  public void testImportCompletionHintForSameDirectoryModuleInPlainDirectory() {
+    doTestVariantTailText("plainDirectory/sample.py", "logging", " (plainDirectory)");
+  }
+
+  // PY-46056
+  public void testFromImportCompletionHintForSameDirectoryModuleInOrdinaryPackage() {
+    doTestVariantTailText("ordinaryPackage/sample.py", "logging", null);
+  }
+
+  // PY-46056
+  public void testFromImportCompletionHintForSameDirectoryModuleInPlainDirectory() {
+    doTestVariantTailText("plainDirectory/sample.py", "logging", " (plainDirectory)");
+  }
+
+  private void doTestVariantTailText(@NotNull String entryFilePath, @NotNull String variantName, @Nullable String tailText) {
+    runWithLanguageLevel(LanguageLevel.getLatest(), () -> {
+      myFixture.copyDirectoryToProject(getTestName(true), "");
+      myFixture.configureByFile(entryFilePath);
+      LookupElement[] variants = myFixture.completeBasic();
+      assertNotNull(variants);
+      LookupElement lookupElement = ContainerUtil.find(variants, v -> v.getLookupString().equals(variantName));
+      assertNotNull(lookupElement);
+      assertEquals(tailText, TestLookupElementPresentation.renderElement(lookupElement).getTailText());
+    });
+  }
+
   // PY-46054
   public void testFromImportFromSameDirectoryModule() {
     runWithLanguageLevel(LanguageLevel.getLatest(), () -> {