[python] PY-82265 edge cases for GTD

(cherry picked from commit 6de5053ef51534f6efe8a1ca4cf428b8d8c44cfa)

IJ-MR-166989

GitOrigin-RevId: 32b61d838cb740f75b4d16842a11bda770ec87fa

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyGotoDeclarationHandler.kt

@@ -21,7 +21,6 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.impl.source.resolve.FileContextUtil
 import com.jetbrains.python.PyUserInitiatedResolvableReference
 import com.jetbrains.python.psi.PyElement
-import com.jetbrains.python.psi.PyImportElement
 import com.jetbrains.python.psi.PyReferenceOwner
 import com.jetbrains.python.psi.resolve.PyResolveContext
 import com.jetbrains.python.psi.resolve.PyResolveUtil
@@ -52,28 +51,18 @@ class PyGotoDeclarationHandler : GotoDeclarationHandlerBase() {
     val parent = sourceElement.parent
     val referenceOwner = sourceElement as? PyReferenceOwner ?: parent as? PyReferenceOwner
     if (referenceOwner != null) {
-      val resolved = PyResolveUtil.multiResolveDeclaration(referenceOwner.getReference(context), context)
-      if (resolved != null) {
-        val stubResults =
-          resolved
-            .filter {
-              PyiUtil.isInsideStub(it) && !PyiUtil.isInsideStub(FileContextUtil.getContextFile(sourceElement)!!)
-            }
-            .map {
+      val results = PyResolveUtil.multiResolveDeclaration(referenceOwner.getReference(context), context)
+          .map {
+            if (PyiUtil.isInsideStub(it) && !PyiUtil.isInsideStub(FileContextUtil.getContextFile(sourceElement)!!)) {
               PyiUtil.getOriginalElement(it as PyElement) ?: it
             }
-        if (stubResults.isNotEmpty()) {
-          return stubResults.toTypedArray()
-        }
-
-       val results = resolved
-        .filter { it !== referenceOwner && it !is PyImportElement }
-        .groupBy { it.containingFile }
-        // include all the definitions from the current file, otherwise just the top result
-        .flatMap { if (it.key == sourceElement.containingFile) it.value else listOf(it.value.first()) }
-        if (results.isNotEmpty()) {
-          return results.toTypedArray()
-        }
+            else it
+          }
+          .filter { it !== referenceOwner }
+          .groupBy { it.containingFile }
+          .flatMap { if (it.key == sourceElement.containingFile) it.value else listOf(it.value.first()) }
+      if (results.isNotEmpty()) {
+        return results.toTypedArray()
       }
     }
 

python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java

@@ -521,7 +521,7 @@ public final class PyResolveUtil {
     return rate;
   }
 
-  public static @Nullable List<PsiElement> multiResolveDeclaration(@NotNull PsiReference reference, @NotNull PyResolveContext resolveContext) {
+  public static @NotNull List<@NotNull PsiElement> multiResolveDeclaration(@NotNull PsiReference reference, @NotNull PyResolveContext resolveContext) {
     final PsiElement element = reference.getElement();
 
     final var context = resolveContext.getTypeEvalContext();
@@ -534,7 +534,7 @@ public final class PyResolveUtil {
 
         final var constructor = ContainerUtil.find(
           PyUtil.filterTopPriorityElements(PyCallExpressionHelper.resolveImplicitlyInvokedMethods(type, call, resolveContext)),
-          it -> it instanceof PyPossibleClassMember && ((PyPossibleClassMember)it).getContainingClass() == cls
+          it -> it instanceof PyPossibleClassMember possibleClassMember && possibleClassMember.getContainingClass() == cls
         );
 
         if (constructor != null) {
@@ -544,16 +544,16 @@ public final class PyResolveUtil {
     }
 
     if (reference instanceof PsiPolyVariantReference multiReference) {
-      return Stream.of(multiReference.multiResolve(false)).map(result -> result.getElement()).toList();
+      return PyUtil.multiResolveTopPriority(multiReference);
     }
     final var result = reference.resolve();
-    if (result == null) return null;
+    if (result == null) return List.of();
     return List.of(result);
   }
 
   public static @Nullable PsiElement resolveDeclaration(@NotNull PsiReference reference, @NotNull PyResolveContext resolveContext) {
     final var result = multiResolveDeclaration(reference, resolveContext);
-    if (result == null || result.isEmpty()) return null;
+    if (result.isEmpty()) return null;
     return result.get(0);
   }
 

python/testData/navigation/test_multi_with_pyi/stubbed.py

@@ -0,0 +1,2 @@
+class A:
+  pass

python/testData/navigation/test_multi_with_pyi/stubbed.pyi

@@ -0,0 +1 @@
+class A: ...

python/testSrc/com/jetbrains/python/PyNavigationTest.kt

@@ -228,7 +228,7 @@ class PyNavigationTest : PyTestCase() {
     myFixture.copyDirectoryToProject(getTestName(true), "")
     myFixture.configureByFile("test.py")
     val target = PyGotoDeclarationHandler().getGotoDeclarationTarget(elementAtCaret, myFixture.editor)
-    TestCase.assertNotNull(target)
+    assertNotNull(target)
     assertInstanceOf(target, PyFunction::class.java)
     checkPyNotPyi(target?.containingFile)
   }
@@ -402,6 +402,27 @@ class PyNavigationTest : PyTestCase() {
     }
   }
 
+  fun `test multi with pyi`() {
+    myFixture.copyDirectoryToProject("test_multi_with_pyi", "")
+    val (stubbed, local) = checkMulti(
+      """
+      if bool():
+          from stubbed import A
+      else:
+          class A:
+              pass
+
+      <caret>A
+      """.trimIndent(),
+      2
+    )
+    checkPyNotPyi(stubbed.containingFile)
+    assertEquals("stubbed.py", stubbed.containingFile.name)
+
+    checkPyNotPyi(local.containingFile)
+    assertEquals("test.py", local.containingFile.name)
+  }
+
   private fun doTestGotoDeclarationNavigatesToPyNotPyi() {
     myFixture.copyDirectoryToProject(getTestName(true), "")
     myFixture.configureByFile("test.py")

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

@@ -58,6 +58,7 @@ import com.jetbrains.python.psi.search.PySearchUtilBase;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.sdk.PythonSdkUtil;
+import kotlin.Pair;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.junit.Assert;