IDEA-CR-51700: go to definition/implementation on file should lead to py files instead of pyi (PY-35129, PY-32345)

GitOrigin-RevId: 0390bc3103ff22903e96d1a669cd4a6aca0980e9

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

@@ -19,10 +19,14 @@ import com.intellij.codeInsight.navigation.actions.GotoDeclarationHandlerBase;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiReference;
+import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyUserInitiatedResolvableReference;
+import com.jetbrains.python.psi.PyElement;
 import com.jetbrains.python.psi.PyReferenceOwner;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.pyi.PyiFile;
+import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -55,7 +59,12 @@ public final class PyGotoDeclarationHandler extends GotoDeclarationHandlerBase {
       referenceOwner = (PyReferenceOwner)parent; //Reference expression may be parent of IDENTIFIER
     }
     if (referenceOwner != null) {
-      return referenceOwner.getReference(context).resolve();
+      final PsiElement resolved = referenceOwner.getReference(context).resolve();
+      if (resolved instanceof PyiFile) {
+        final PsiElement original = PyiUtil.getOriginalElement(((PyElement)resolved));
+        return ObjectUtils.chooseNotNull(original, resolved);
+      }
+      return resolved;
     }
     // If element is not ref owner, it still may have provided references, lets find some
     final PsiElement element = findProvidedReferenceAndResolve(sourceElement);

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

@@ -241,6 +241,7 @@ private fun resultsFromRoots(name: QualifiedName, context: PyQualifiedNameResolv
   val sdk = context.effectiveSdk
   val module = context.module
   val footholdFile = context.footholdFile
+  val withoutStubs = context.withoutStubs
 
   val visitor = RootVisitor { root, module, sdk, isModuleSource ->
     val results = if (isModuleSource) moduleResults else sdkResults
@@ -250,6 +251,9 @@ private fun resultsFromRoots(name: QualifiedName, context: PyQualifiedNameResolv
         effectiveSdk != null && PyTypeShed.isInside(root) && !PyTypeShed.maySearchForStubInRoot(name, root, effectiveSdk)) {
       return@RootVisitor true
     }
+    if (withoutStubs && (PyTypeShed.isInside(root) || isInStubPackage(PsiManager.getInstance(context.project).findDirectory(root)!!))) {
+        return@RootVisitor true
+    }
     results.addAll(resolveInRoot(name, root, context))
     if (isAcceptRootAsTopLevelPackage(context) && name.matchesPrefix(
         QualifiedName.fromDottedString(root.name))) {

python/python-psi-impl/src/com/jetbrains/python/psi/search/PyDefinitionsSearch.java

@@ -10,6 +10,8 @@ import com.jetbrains.python.psi.PyAssignmentStatement;
 import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyFunction;
 import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.pyi.PyiFile;
+import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -35,6 +37,12 @@ public class PyDefinitionsSearch implements QueryExecutor<PsiElement, PsiElement
         return consumer.process(parent);
       }
     }
+    else if (queryParameters instanceof PyiFile) {
+      final PsiElement originalElement = ReadAction.compute(() -> PyiUtil.getOriginalElement((PyiFile)queryParameters));
+      if (originalElement != null) {
+        consumer.process(originalElement);
+      }
+    }
     return true;
   }
 }

python/testData/navigation/goToDeclarationForDirectory/test.py

@@ -0,0 +1 @@
+import collection<caret>s

python/testData/navigation/onPyiFile/source.py

@@ -0,0 +1,2 @@
+def foo():
+    return "bar"

python/testData/navigation/onPyiFile/source.pyi

@@ -0,0 +1 @@
+def foo() -> str: ...

python/testData/navigation/onPyiFile/test.py

@@ -0,0 +1 @@
+from sou<caret>rce import foo

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

@@ -0,0 +1,54 @@
+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python
+
+import com.intellij.psi.PsiElement
+import com.intellij.testFramework.LightProjectDescriptor
+import com.intellij.testFramework.fixtures.CodeInsightTestUtil
+import com.jetbrains.python.fixtures.PyTestCase
+import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.psi.impl.PyGotoDeclarationHandler
+import com.jetbrains.python.pyi.PyiFile
+
+class PyNavigationTest : PyTestCase() {
+
+  // PY-35129
+  fun testGoToDeclarationOnPyiFile() {
+    configureByDir("onPyiFile")
+    val target = PyGotoDeclarationHandler().getGotoDeclarationTarget(elementAtCaret, myFixture.editor)
+    checkPyNotPyi(target)
+  }
+
+  // PY-35129
+  fun testGoToImplementationOnPyiFile() {
+    configureByDir("onPyiFile")
+    val gotoData = CodeInsightTestUtil.gotoImplementation(myFixture.editor, myFixture.file)
+    assertSize(1, gotoData.targets)
+    checkPyNotPyi(gotoData.targets[0])
+  }
+
+  // PY-35129
+  fun testGoToDeclarationForDirectory() {
+    configureByDir(getTestName(true))
+    val target = PyGotoDeclarationHandler().getGotoDeclarationTarget(elementAtCaret, myFixture.editor)
+    checkPyNotPyi(target)
+  }
+
+  private fun configureByDir(dirName: String) {
+    myFixture.copyDirectoryToProject(dirName, "")
+    myFixture.configureByFile("test.py")
+    assertTrue(myFixture.elementAtCaret is PyiFile)
+  }
+
+  private fun checkPyNotPyi(file: PsiElement?) {
+    assertTrue(file is PyFile)
+    assertTrue(file !is PyiFile)
+  }
+
+  override fun getTestDataPath(): String {
+    return super.getTestDataPath() + "/navigation"
+  }
+
+  override fun getProjectDescriptor(): LightProjectDescriptor? = ourPy3Descriptor
+
+
+}