From bd2c7616a2847d30a28a55920448b5da42eb2380 Mon Sep 17 00:00:00 2001
From: Morgan Bartholomew <morgan.bartholomew@jetbrains.com>
Date: Fri, 31 Oct 2025 22:53:57 +1000
Subject: [PATCH] [python] PY-85012 fix `PyGotoDeclarationHandler` for multi
 assigned local variables

(cherry picked from commit 27ae2bddd6c2b0b17a8bd2ace3c8da0229afc710)

IJ-MR-180935

GitOrigin-RevId: 02d9e5f80813fb326d0fa0c9aec38e0caec32bb5
---
 .../python/psi/impl/PyGotoDeclarationHandler.kt   | 13 ++++++++++++-
 .../com/jetbrains/python/PyNavigationTest.kt      | 15 +++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

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 6b099f6789be..9be363b93a09 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
@@ -18,10 +18,13 @@ package com.jetbrains.python.psi.impl
 import com.intellij.codeInsight.navigation.actions.GotoDeclarationHandlerBase
 import com.intellij.openapi.editor.Editor
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
 import com.intellij.psi.impl.source.resolve.FileContextUtil
 import com.jetbrains.python.PyUserInitiatedResolvableReference
 import com.jetbrains.python.psi.PyElement
+import com.jetbrains.python.psi.PyQualifiedExpression
 import com.jetbrains.python.psi.PyReferenceOwner
+import com.jetbrains.python.psi.PyTargetExpression
 import com.jetbrains.python.psi.resolve.PyResolveContext
 import com.jetbrains.python.psi.resolve.PyResolveUtil
 import com.jetbrains.python.psi.types.TypeEvalContext
@@ -60,7 +63,15 @@ class PyGotoDeclarationHandler : GotoDeclarationHandlerBase() {
           }
           .filter { it !== referenceOwner }
           .groupBy { it.containingFile }
-          .flatMap { if (it.key == sourceElement.containingFile) it.value else listOf(it.value.first()) }
+          .flatMap { (containingFile, declarations) ->
+            if (containingFile != sourceElement.containingFile)
+              declarations.take(1)
+              // if it's a qualified expression, then it could be a union, so go to the different declarations
+              //  otherwise go to the most recent assignment
+              else if ((referenceOwner as? PyQualifiedExpression)?.isQualified == true)
+                declarations
+                else declarations.takeLast(1)
+          }
       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 25dfadfc9861..f1f45bea8b48 100644
--- a/python/testSrc/com/jetbrains/python/PyNavigationTest.kt
+++ b/python/testSrc/com/jetbrains/python/PyNavigationTest.kt
@@ -4,6 +4,7 @@ package com.jetbrains.python
 import com.intellij.codeInsight.navigation.actions.GotoDeclarationOrUsageHandler2
 import com.intellij.codeInsight.navigation.actions.GotoDeclarationOrUsageHandler2.GTDUOutcome
 import com.intellij.ide.util.gotoByName.GotoSymbolModel2
+import com.intellij.idea.TestFor
 import com.intellij.psi.PsiElement
 import com.intellij.psi.util.parentOfType
 import com.intellij.testFramework.fixtures.CodeInsightTestUtil
@@ -423,6 +424,20 @@ class PyNavigationTest : PyTestCase() {
     assertEquals("test.py", local.containingFile.name)
   }
 
+  @TestFor(issues = ["PY-85012"])
+  fun `test gtd for local variable`() {
+    val (result) = checkMulti(
+      """
+      a = 1
+      if bool():
+          a = 2
+      <caret>a
+    """.trimIndent(),
+      1,
+    )
+    assertEquals("a = 2", result.parent.text)
+  }
+
   private fun doTestGotoDeclarationNavigatesToPyNotPyi() {
     myFixture.copyDirectoryToProject(getTestName(true), "")
     myFixture.configureByFile("test.py")