PY-38636 Fix go to declaration lead to pyi in some cases

(cherry picked from commit 33b52287f1d5241ee0bc4b8ff1c065e32f26d063)

IJ-MR-18958

GitOrigin-RevId: 5a59d6e6ea3ab001b8237d408340536abe2678e0

python/python-psi-impl/src/com/jetbrains/python/codeInsight/PyTargetElementEvaluator.java

@@ -15,24 +15,26 @@
  */
 package com.jetbrains.python.codeInsight;
 
-import com.intellij.codeInsight.TargetElementEvaluatorEx2;
+import com.intellij.codeInsight.TargetElementEvaluator;
 import com.intellij.codeInsight.TargetElementUtilBase;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
-import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.PyParameter;
+import com.jetbrains.python.psi.PyReferenceExpression;
+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;
-import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.HashSet;
 import java.util.Set;
 
-public class PyTargetElementEvaluator extends TargetElementEvaluatorEx2 {
+public class PyTargetElementEvaluator implements TargetElementEvaluator {
   @Override
   public boolean includeSelfInGotoImplementation(@NotNull PsiElement element) {
     return false;
@@ -66,14 +68,4 @@ public class PyTargetElementEvaluator extends TargetElementEvaluatorEx2 {
     }
     return result;
   }
-
-  @Nullable
-  @Override
-  public PsiElement getGotoDeclarationTarget(@NotNull final PsiElement element, @Nullable final PsiElement navElement) {
-    if (element instanceof PyElement) {
-      final PsiElement originalElement = PyiUtil.getOriginalElement((PyElement)element);
-      return originalElement != null ? originalElement : navElement;
-    }
-    return navElement;
-  }
 }

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

@@ -21,8 +21,7 @@ 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.*;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import com.jetbrains.python.psi.types.TypeEvalContext;
@@ -31,6 +30,8 @@ import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Objects;
+
 /**
  * {@link com.intellij.codeInsight.navigation.actions.GotoDeclarationAction} uses {@link PsiElement#findReferenceAt(int)}.
  * This method knows nothing about execution context and {@link PyBaseElementImpl} injects loose {@link TypeEvalContext}.
@@ -60,10 +61,9 @@ public final class PyGotoDeclarationHandler extends GotoDeclarationHandlerBase {
       referenceOwner = (PyReferenceOwner)parent; //Reference expression may be parent of IDENTIFIER
     }
     if (referenceOwner != null) {
-      final PsiElement resolved = PyResolveUtil.resolveDeclaration(referenceOwner.getReference(context), context);
-      if (resolved instanceof PyiFile) {
-        final PsiElement original = PyiUtil.getOriginalElement(((PyElement)resolved));
-        return ObjectUtils.chooseNotNull(original, resolved);
+      PsiElement resolved = PyResolveUtil.resolveDeclaration(referenceOwner.getReference(context), context);
+      if (resolved != null && resolved.getContainingFile() instanceof PyiFile) {
+        return ObjectUtils.chooseNotNull(PyiUtil.getOriginalElement((PyElement)resolved), resolved);
       }
       return resolved != referenceOwner ? resolved : null;
     }

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

@@ -26,6 +26,7 @@ import com.jetbrains.python.psi.types.PyModuleType;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.pyi.PyiStubSuppressor;
+import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -392,7 +393,9 @@ public final class ResolveImportUtil {
 
     final PsiDirectory subdir = dir.findSubdirectory(referencedName);
     // VFS may be case insensitive on Windows, but resolve is always case sensitive (PEP 235, PY-18958), so we check name here
-    if (subdir != null && subdir.getName().equals(referencedName) && (!checkForPackage || PyUtil.isPackage(subdir, containingFile))) {
+    if (subdir != null && subdir.getName().equals(referencedName) &&
+        (!checkForPackage || PyUtil.isPackage(subdir, containingFile)) &&
+        (!withoutStubs || !PyiUtil.isPyiFileOfPackage(subdir))) {
       result.add(new RatedResolveResult(RatedResolveResult.RATE_NORMAL, PyStubPackages.transferStubPackageMarker(dir, subdir)));
     }
 

python/python-psi-impl/src/com/jetbrains/python/pyi/PyiUtil.java

@@ -67,13 +67,20 @@ public final class PyiUtil {
   @Nullable
   public static PsiElement getOriginalElement(@NotNull PyElement element) {
     final PsiFile file = element.getContainingFile();
-    if (file instanceof PyiFile) {
-      final PyFile originalFile = getOriginalFile((PyiFile)file);
-      if (originalFile != null) {
-        return findSimilarElement(element, originalFile);
+    if (!(file instanceof PyiFile)) return null;
+
+    final PyFile originalFile = getOriginalFile((PyiFile)file);
+    if (originalFile == null) return null;
+
+    PsiElement result = findSimilarElement(element, originalFile);
+    if (result == null && element instanceof PyFunction) {
+      PyClass containingClass = PyUtil.turnConstructorIntoClass((PyFunction)element);
+      if (containingClass != null) {
+        result = findSimilarElement(containingClass, originalFile);
       }
     }
-    return null;
+
+    return result;
   }
 
   /**

python/testData/navigation/classInPyiAssignedToFunctionInPy/source.py

@@ -0,0 +1,4 @@
+def foo():
+    return 42
+
+Clazz = foo()

python/testData/navigation/classInPyiAssignedToFunctionInPy/source.pyi

@@ -0,0 +1,2 @@
+class Clazz:
+    def __init__(self) -> None: ...

python/testData/navigation/classInPyiAssignedToFunctionInPy/test.py

@@ -0,0 +1,2 @@
+import source
+cl = source.Cl<caret>azz()

python/testData/navigation/classInPyiClassInPy/source.py

@@ -0,0 +1,3 @@
+class Clazz:
+    def __init__(self) -> None:
+        pass

python/testData/navigation/classInPyiClassInPy/source.pyi

@@ -0,0 +1,2 @@
+class Clazz:
+    def __init__(self) -> None: ...

python/testData/navigation/classInPyiClassInPy/test.py

@@ -0,0 +1,2 @@
+import source
+cl = source.Cl<caret>azz()

python/testData/navigation/importFile/addRoots/mypack/source.py

@@ -0,0 +1,3 @@
+class Clazz:
+    def __init__(self) -> None:
+        pass

python/testData/navigation/importFile/mypack/source.pyi

@@ -0,0 +1,2 @@
+class Clazz:
+    def __init__(self) -> None: ...

python/testData/navigation/importFile/test.py

@@ -0,0 +1,3 @@
+from mypack.sour<caret>ce import Clazz
+
+c = Clazz()

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

@@ -16,6 +16,7 @@ import com.jetbrains.python.psi.impl.PyPsiUtils
 import com.jetbrains.python.psi.types.TypeEvalContext
 import com.jetbrains.python.pyi.PyiFile
 import com.jetbrains.python.pyi.PyiUtil
+import junit.framework.TestCase
 
 class PyNavigationTest : PyTestCase() {
 
@@ -187,6 +188,36 @@ class PyNavigationTest : PyTestCase() {
     assertFalse(PyiUtil.isOverload(foo, context))
   }
 
+  // PY-38636
+  fun testClassInPyiAssignedToFunctionInPy() {
+    myFixture.copyDirectoryToProject(getTestName(true), "")
+    myFixture.configureByFile("test.py")
+    val target = PyGotoDeclarationHandler().getGotoDeclarationTarget(elementAtCaret, myFixture.editor)
+    TestCase.assertNotNull(target)
+    assertInstanceOf(target, PyTargetExpression::class.java)
+    checkPyNotPyi(target?.containingFile)
+  }
+
+  // PY-38636
+  fun testStubInUserCode() {
+    myFixture.copyDirectoryToProject("importFile", "")
+    myFixture.configureByFile("test.py")
+    runWithAdditionalClassEntryInSdkRoots(myFixture.findFileInTempDir("addRoots")) {
+      val target = PyGotoDeclarationHandler().getGotoDeclarationTarget(elementAtCaret, myFixture.editor)
+      checkPyNotPyi(target?.containingFile)
+    }
+  }
+
+  // PY-38636
+  fun testClassInPyiClassInPy() {
+    myFixture.copyDirectoryToProject(getTestName(true), "")
+    myFixture.configureByFile("test.py")
+    val target = PyGotoDeclarationHandler().getGotoDeclarationTarget(elementAtCaret, myFixture.editor)
+    TestCase.assertNotNull(target)
+    assertInstanceOf(target, PyFunction::class.java)
+    checkPyNotPyi(target?.containingFile)
+  }
+
   private fun doTestGotoDeclarationOrUsagesOutcome(expectedOutcome: GTDUOutcome, text: String) {
     myFixture.configureByText("a.py", text)
     val actualOutcome = GotoDeclarationOrUsageHandler2.testGTDUOutcome(myFixture.editor, myFixture.file, myFixture.caretOffset)