From 09659aec5b2893fbc5cebb7aef69df76bfd0f080 Mon Sep 17 00:00:00 2001
From: Morgan Bartholomew <morgan.bartholomew@jetbrains.com>
Date: Wed, 11 Jun 2025 00:10:19 +1000
Subject: [PATCH] [python] PY-81789 filter unhelpful declarations for GTD

(cherry picked from commit 0fba00400382b32bbdbb1d6112126b4097895f2f)

IJ-MR-165350

GitOrigin-RevId: af95c00431a09faf9a08bb3c8a7f88b8d3376dee
---
 .../psi/impl/PyGotoDeclarationHandler.kt      |  7 +++-
 .../com/jetbrains/python/PyNavigationTest.kt  | 40 ++++++++++++++-----
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyGotoDeclarationHandler.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyGotoDeclarationHandler.kt
index 9bf7004f1697..597a2eaf248e 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyGotoDeclarationHandler.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyGotoDeclarationHandler.kt
@@ -21,6 +21,7 @@ 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
@@ -65,7 +66,11 @@ class PyGotoDeclarationHandler : GotoDeclarationHandlerBase() {
           return stubResults.toTypedArray()
         }
 
-        val results = resolved.filter { it !== referenceOwner }
+       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()
         }
diff --git a/python/testSrc/com/jetbrains/python/PyNavigationTest.kt b/python/testSrc/com/jetbrains/python/PyNavigationTest.kt
index 9d0726d3cd95..15dd973770b9 100644
--- a/python/testSrc/com/jetbrains/python/PyNavigationTest.kt
+++ b/python/testSrc/com/jetbrains/python/PyNavigationTest.kt
@@ -345,10 +345,9 @@ class PyNavigationTest : PyTestCase() {
             x: int
         ab: A | B
         ab.<caret>x
-        """.trimIndent()
-    ).also {
-      assertSize(2, it)
-    }
+        """.trimIndent(),
+      2,
+    )
 
     assertInstanceOf<PyTargetExpression>(a)
     assertEquals("A", a.parentOfType<PyClass>()!!.name)
@@ -368,10 +367,9 @@ class PyNavigationTest : PyTestCase() {
             def __add__(self, other): ...
             def __radd__(self, other): ...
         A() +<caret> B()
-      """.trimIndent()
-    ).also {
-      assertSize(2, it)
-    }
+      """.trimIndent(),
+      2,
+    )
 
     assertInstanceOf<PyFunction>(a)
     assertEquals("__add__", a.name)
@@ -383,6 +381,27 @@ class PyNavigationTest : PyTestCase() {
 
   }
 
+  // PY-81789
+  fun `test import isn't declaration and only one from other file`() {
+    runWithAdditionalFileInLibDir(
+      "a.py",
+      """
+      a = 1
+      a = 2""".trimIndent()
+    ) {
+      val (a) = checkMulti(
+        """
+        from a import a
+        
+        <caret>a
+        """.trimIndent(),
+        1
+      )
+
+      assertInstanceOf<PyTargetExpression>(a)
+    }
+  }
+
   private fun doTestGotoDeclarationNavigatesToPyNotPyi() {
     myFixture.copyDirectoryToProject(getTestName(true), "")
     myFixture.configureByFile("test.py")
@@ -414,11 +433,14 @@ class PyNavigationTest : PyTestCase() {
   }
 
 
-  private fun checkMulti(text: String): List<PsiElement> {
+  private fun checkMulti(text: String, expectedSize: Int?): List<PsiElement> {
     myFixture.configureByText("test.py", text)
     return PyGotoDeclarationHandler()
       .getGotoDeclarationTargets(elementAtCaret, -1, myFixture.editor)!!
       .toList()
+      .also {
+        if (expectedSize != null) assertSize(expectedSize, it)
+      }
   }
 
   private fun checkPyNotPyi(file: PsiElement?) {