From f083128e86f3bd33897b6789d478fd6fe738ed0c Mon Sep 17 00:00:00 2001 From: Morgan Bartholomew Date: Wed, 25 Jun 2025 17:33:35 +0200 Subject: [PATCH] [python] PY-82265 edge cases for GTD GitOrigin-RevId: 6de5053ef51534f6efe8a1ca4cf428b8d8c44cfa --- .../psi/impl/PyGotoDeclarationHandler.kt | 31 ++++++------------- .../python/psi/resolve/PyResolveUtil.java | 10 +++--- .../navigation/test_multi_with_pyi/stubbed.py | 2 ++ .../test_multi_with_pyi/stubbed.pyi | 1 + .../com/jetbrains/python/PyNavigationTest.kt | 23 +++++++++++++- .../jetbrains/python/fixtures/PyTestCase.java | 1 + 6 files changed, 41 insertions(+), 27 deletions(-) create mode 100644 python/testData/navigation/test_multi_with_pyi/stubbed.py create mode 100644 python/testData/navigation/test_multi_with_pyi/stubbed.pyi 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 597a2eaf248e..6b099f6789be 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,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() } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java index 32dd1cac2e27..b530caf4cc39 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java @@ -515,7 +515,7 @@ public final class PyResolveUtil { return rate; } - public static @Nullable List 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(); @@ -528,7 +528,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) { @@ -538,16 +538,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); } diff --git a/python/testData/navigation/test_multi_with_pyi/stubbed.py b/python/testData/navigation/test_multi_with_pyi/stubbed.py new file mode 100644 index 000000000000..d248dc1809dd --- /dev/null +++ b/python/testData/navigation/test_multi_with_pyi/stubbed.py @@ -0,0 +1,2 @@ +class A: + pass \ No newline at end of file diff --git a/python/testData/navigation/test_multi_with_pyi/stubbed.pyi b/python/testData/navigation/test_multi_with_pyi/stubbed.pyi new file mode 100644 index 000000000000..2335dbbf7fa2 --- /dev/null +++ b/python/testData/navigation/test_multi_with_pyi/stubbed.pyi @@ -0,0 +1 @@ +class A: ... \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyNavigationTest.kt b/python/testSrc/com/jetbrains/python/PyNavigationTest.kt index 15dd973770b9..25dfadfc9861 100644 --- a/python/testSrc/com/jetbrains/python/PyNavigationTest.kt +++ b/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 + + 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") diff --git a/python/testSrc/com/jetbrains/python/fixtures/PyTestCase.java b/python/testSrc/com/jetbrains/python/fixtures/PyTestCase.java index c27b842de314..fdeb8adab833 100644 --- a/python/testSrc/com/jetbrains/python/fixtures/PyTestCase.java +++ b/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;