PY-45776 EA-248324 Fix npe in isSameDirectoryResult

GitOrigin-RevId: e3c26a84d89bebe99443033a642d655c1e518c82

platform/core-api/src/com/intellij/psi/util/QualifiedName.java

@@ -104,7 +104,7 @@ public final class QualifiedName implements Comparable<QualifiedName> {
     return true;
   }
 
-  public boolean matchesPrefix(QualifiedName prefix) {
+  public boolean matchesPrefix(@NotNull QualifiedName prefix) {
     if (getComponentCount() < prefix.getComponentCount()) {
       return false;
     }

python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt

@@ -14,6 +14,7 @@ import com.intellij.openapi.roots.FileIndexFacade
 import com.intellij.openapi.util.Ref
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.*
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.QualifiedName
 import com.jetbrains.python.codeInsight.typing.PyTypeShed
 import com.jetbrains.python.codeInsight.typing.isInInlinePackage
@@ -338,8 +339,15 @@ private fun findCache(context: PyQualifiedNameResolveContext): PythonPathCache?
 private fun isSameDirectoryResult(element: PsiElement, context: PyQualifiedNameResolveContext, name: QualifiedName): Boolean {
   if (context.relativeLevel != 0) return false
   val sameDirectoryImportsEnabled = !ResolveImportUtil.isAbsoluteImportEnabledFor(context.foothold)
-  return sameDirectoryImportsEnabled && element is PsiFileSystemItem &&
-         !name.matchesPrefix(QualifiedNameFinder.findShortestImportableQName(element))
+  if (!sameDirectoryImportsEnabled || element !is PsiFileSystemItem) return false
+  val shortestImportableQName = QualifiedNameFinder.findShortestImportableQName(element)
+  if (shortestImportableQName != null) {
+    return name != shortestImportableQName
+  }
+  else {
+    val footholdDir = context.containingDirectory ?: return false
+    return PsiTreeUtil.isAncestor(footholdDir, element, true)
+  }
 }
 
 private fun isSameDirectoryOrRelativeImportResult(name: QualifiedName, directory: PsiDirectory, result: PsiElement,

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/ordinaryPackageImportPrioritizeExcludedDirectoryInRootOverSameDirectoryModule/ordinaryPackage/excluded.py

@@ -0,0 +1,2 @@
+def foo():
+    pass

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/ordinaryPackageImportPrioritizeExcludedDirectoryInRootOverSameDirectoryModule/ordinaryPackage/script.py

@@ -0,0 +1,2 @@
+from excluded import foo
+#       <ref>

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/ordinaryPackageImportPrioritizeModuleInRootOverSameDirectoryExcludedDirectory/excluded.py

@@ -0,0 +1,2 @@
+def foo():
+    pass

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/ordinaryPackageImportPrioritizeModuleInRootOverSameDirectoryExcludedDirectory/ordinaryPackage/script.py

@@ -0,0 +1,2 @@
+from excluded import foo
+#       <ref>

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/ordinaryPackageInvalidNameImportPrioritizeModuleInRootOverSameDirectoryModule/ordinary package/main.py

@@ -0,0 +1,2 @@
+import mod
+#       <ref>

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/plainDirectoryImportResolveExcludedDirectoryModuleNotThrowsException/plainDirectory/excluded.py

@@ -0,0 +1,2 @@
+def foo():
+    pass

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/plainDirectoryImportResolveExcludedDirectoryModuleNotThrowsException/plainDirectory/script.py

@@ -0,0 +1,2 @@
+from excluded import foo
+#       <ref>

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/plainDirectoryImportResolveSameDirectoryModuleNotThrowsException/not-valid-identifier/script.py

@@ -0,0 +1,2 @@
+import lib
+#       <ref>

python/testData/resolve/multiFile/relativeAndSameDirectoryImports/plainDirectoryInvalidNameImportPrioritizeSameDirectoryModuleOverModuleInRoot/plain directory/main.py

@@ -0,0 +1,2 @@
+import mod
+#       <ref>

python/testSrc/com/jetbrains/python/PyRelativeImportResolveTest.java

@@ -19,6 +19,7 @@ import com.intellij.application.options.RegistryManager;
 import com.intellij.psi.*;
 import com.intellij.psi.util.QualifiedName;
 import com.intellij.testFramework.LightProjectDescriptor;
+import com.intellij.testFramework.PsiTestUtil;
 import com.jetbrains.python.fixtures.PyMultiFileResolveTestCase;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.namespacePackages.PyNamespacePackagesService;
@@ -346,20 +347,73 @@ public class PyRelativeImportResolveTest extends PyMultiFileResolveTestCase {
     });
   }
 
-  private <T extends PsiElement> void assertResolvesInsideProjectTo(@NotNull Class<T> cls) {
-    PsiManager psiManager = myFixture.getPsiManager();
-    assertTrue(psiManager.isInProject(assertResolvesToElementOfClass(cls)));
+  // PY-45776
+  public void testPlainDirectoryImportResolveSameDirectoryModuleNotThrowsException() {
+    myTestFileName = "not-valid-identifier/script.py";
+    assertResolvesInsideProjectTo(PyFile.class);
   }
 
-  private <T extends PsiElement> void assertResolvesOutsideProjectTo(@NotNull Class<T> cls) {
+  // PY-45776
+  public void testPlainDirectoryImportResolveExcludedDirectoryModuleNotThrowsException() {
+    myTestFileName = PLAIN_DIR + "/script.py";
+    PsiFile currentFile = prepareFile();
+    PsiTestUtil.addExcludedRoot(myFixture.getModule(), myFixture.findFileInTempDir("excluded"));
+    PsiElement element = doResolve(currentFile);
+    assertInstanceOf(element, PyFile.class);
+    assertEquals(myFixture.findFileInTempDir(PLAIN_DIR + "/excluded.py"), ((PyFile)element).getVirtualFile());
+  }
+
+  // PY-45776
+  public void testOrdinaryPackageInvalidNameImportPrioritizeModuleInRootOverSameDirectoryModule() {
+    myTestFileName = "ordinary package/main.py";
+    PsiFile file = assertResolvesInsideProjectTo(PyFile.class);
+    assertEquals(myFixture.findFileInTempDir("mod.py"), file.getVirtualFile());
+  }
+
+  // PY-45776
+  public void testPlainDirectoryInvalidNameImportPrioritizeSameDirectoryModuleOverModuleInRoot() {
+    myTestFileName = "plain directory/main.py";
+    PsiFile file = assertResolvesInsideProjectTo(PyFile.class);
+    assertEquals(myFixture.findFileInTempDir("plain directory/mod.py"), file.getVirtualFile());
+  }
+
+  // PY-45776
+  public void testOrdinaryPackageImportPrioritizeExcludedDirectoryInRootOverSameDirectoryModule() {
+    myTestFileName = ORDINARY_PACK_DIR + "/script.py";
+    PsiFile currentFile = prepareFile();
+    PsiTestUtil.addExcludedRoot(myFixture.getModule(), myFixture.findFileInTempDir("excluded"));
+    PsiElement element = doResolve(currentFile);
+    assertInstanceOf(element, PsiDirectory.class);
+    assertEquals(myFixture.findFileInTempDir("excluded"), ((PsiDirectory)element).getVirtualFile());
+  }
+
+  // PY-45776
+  public void testOrdinaryPackageImportPrioritizeModuleInRootOverSameDirectoryExcludedDirectory() {
+    myTestFileName = ORDINARY_PACK_DIR + "/script.py";
+    PsiFile currentFile = prepareFile();
+    PsiTestUtil.addExcludedRoot(myFixture.getModule(), myFixture.findFileInTempDir(ORDINARY_PACK_DIR + "/excluded"));
+    PsiElement element = doResolve(currentFile);
+    assertInstanceOf(element, PyFile.class);
+    assertEquals(myFixture.findFileInTempDir("excluded.py"), ((PyFile)element).getVirtualFile());
+  }
+
+  private <T extends PsiElement> T assertResolvesInsideProjectTo(@NotNull Class<T> cls) {
     PsiManager psiManager = myFixture.getPsiManager();
-    assertFalse(psiManager.isInProject(assertResolvesToElementOfClass(cls)));
+    T resolved = assertResolvesToElementOfClass(cls);
+    assertTrue(psiManager.isInProject(resolved));
+    return resolved;
+  }
+
+  private <T extends PsiElement> T assertResolvesOutsideProjectTo(@NotNull Class<T> cls) {
+    PsiManager psiManager = myFixture.getPsiManager();
+    T resolved = assertResolvesToElementOfClass(cls);
+    assertFalse(psiManager.isInProject(resolved));
+    return resolved;
   }
 
   private <T extends PsiElement> @NotNull T assertResolvesToElementOfClass(@NotNull Class<T> cls) {
     PsiElement resolved = doResolve();
-    assertInstanceOf(resolved, cls);
-    return (T) resolved;
+    return assertInstanceOf(resolved, cls);
   }
 
   private void toggleNamespacePackageDirectory(@NotNull String directory) {