PY-54905 PY-54620 PY-61740 Fix goto declaration and implementation resolved pyi not py

PyDefinitionsSearch.java: replace psi element on original element if it is from pyi and then search for overrides, inheritors etc. PyiUtil.java: recursively get original element while it is in stub (see the comment) and tests testGoToImplementationNameReExportedThroughAssignmentInPyiStub, testGoToImplementationNameReExportedThroughAssignmentInPyiStubTwice PyUtil.java: make different methods turnDirIntoInitPy, turnDirIntoInitPyi instead of only turnDirIntoInit to handle cases when the same directory contains __init__.pyi and __init__.py files (PY-54620) GitOrigin-RevId: 8f0d8a8133548e1a9f52f93b42aa9cae2302e8d7
2026-09-27 10:03:11 +07:00 · 2023-07-24 14:30:11 +00:00
parent 9aab693d8a
commit 213b82bde3
57 changed files with 261 additions and 32 deletions
@@ -913,6 +913,22 @@ public final class PyUtil {
    } // don't touch non-dirs
  }

+  @Nullable
+  public static PsiElement turnDirIntoInitPy(@Nullable PsiElement target) {
+    if (!(target instanceof PsiDirectory psiDirectory)) return target;
+    return psiDirectory.findFile(PyNames.INIT_DOT_PY);
+  }
+
+  @Nullable
+  public static PsiElement turnDirIntoInitPyi(@Nullable PsiElement target) {
+    if (!(target instanceof PsiDirectory psiDirectory)) return target;
+    final PsiFile initStub = psiDirectory.findFile(PyNames.INIT_DOT_PYI);
+    if (initStub != null && !PyiStubSuppressor.isIgnoredStub(initStub)) {
+      return initStub;
+    }
+    return null;
+  }
+
  /**
   * If directory is a PsiDirectory, that is also a valid Python package, return PsiFile that points to __init__.py,
   * if such file exists, or directory itself (i.e. namespace package). Otherwise, return {@code null}.
@@ -388,7 +388,7 @@ public final class ResolveImportUtil {
    // 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)) &&
-        (!withoutStubs || !PyiUtil.isPyiFileOfPackage(subdir))) {
+        (!withoutStubs || PyUtil.isOrdinaryPackage(subdir) || !PyiUtil.isPyiFileOfPackage(subdir))) {
      result.add(new RatedResolveResult(RatedResolveResult.RATE_NORMAL, PyStubPackages.transferStubPackageMarker(dir, subdir)));
    }

@@ -4,43 +4,50 @@ package com.jetbrains.python.psi.search;
 import com.intellij.openapi.application.ReadAction;
 import com.intellij.psi.PsiElement;
 import com.intellij.util.Processor;
-import com.intellij.util.Query;
 import com.intellij.util.QueryExecutor;
-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.psi.*;
 import com.jetbrains.python.pyi.PyiUtil;
 import org.jetbrains.annotations.NotNull;


 public class PyDefinitionsSearch implements QueryExecutor<PsiElement, PsiElement> {
  @Override
-  public boolean execute(@NotNull final PsiElement e, @NotNull final Processor<? super PsiElement> consumer) {
-    if (e instanceof PyClass) {
-      final Query<PyClass> query = PyClassInheritorsSearch.search((PyClass)e, true);
-      return query.forEach(consumer);
-    }
-    else if (e instanceof PyFunction) {
-      final Query<PyFunction> query =
-        ReadAction.compute(() -> PyOverridingMethodsSearch.search((PyFunction)e, true));
-
-      return query.forEach(consumer);
-    }
-    else if (e instanceof PyTargetExpression) {  // PY-237
-      final PsiElement parent = ReadAction.compute(() -> e.getParent());
-
-      if (parent instanceof PyAssignmentStatement) {
-        return consumer.process(parent);
+  public boolean execute(@NotNull PsiElement element, @NotNull Processor<? super PsiElement> consumer) {
+    if (element instanceof PyElement) {
+      PsiElement finalElement = element;
+      boolean isInsideStub = ReadAction.compute(() -> PyiUtil.isInsideStub(finalElement));
+      if (isInsideStub) {
+        var originalElement = ReadAction.compute(() -> PyiUtil.getOriginalElement((PyElement)finalElement));
+        if (originalElement != null) {
+          element = originalElement;
+          if (!consumer.process(element)) return false;
+        }
      }
    }
-    else if (e instanceof PyiFile) {
-      final PsiElement originalElement = ReadAction.compute(() -> PyiUtil.getOriginalElement((PyiFile)e));
-      if (originalElement != null) {
-        consumer.process(originalElement);
-      }
+    if (element instanceof PyClass) {
+      return processInheritors((PyClass)element, consumer);
+    }
+    else if (element instanceof PyFunction) {
+      return processOverridingMethods((PyFunction)element, consumer);
+    }
+    else if (element instanceof PyTargetExpression) {
+      return processAssignmentStatement(element, consumer);
    }
    return true;
  }
+
+  private static boolean processInheritors(@NotNull PyClass pyClass, @NotNull Processor<? super PsiElement> consumer) {
+    return ReadAction.compute(() -> PyClassInheritorsSearch.search(pyClass, true)).forEach(consumer);
+  }
+
+  private static boolean processOverridingMethods(@NotNull PyFunction pyFunction,
+                                                  @NotNull Processor<? super PsiElement> consumer) {
+    return ReadAction.compute(() -> PyOverridingMethodsSearch.search(pyFunction, true)).forEach(consumer);
+  }
+
+  private static boolean processAssignmentStatement(@NotNull PsiElement element,
+                                                    @NotNull Processor<? super PsiElement> consumer) {
+    PsiElement parent = ReadAction.compute(() -> element.getParent());
+    return !(parent instanceof PyAssignmentStatement) || consumer.process(parent);
+  }
 }
@@ -70,7 +70,20 @@ public final class PyiUtil {
    if (originalFile == null) return null;

    PsiElement result = findSimilarElement(element, originalFile);
-    if (result == null && element instanceof PyFunction) {
+    
+    // If a name is defined in a .pyi stub and the corresponding .py module in a different manner, e.g.
+    // it's exported through an assignment to a top-level attribute in a .pyi stub, but though a regular
+    // "from" import in .py file, we might end up in another stub file again. It happens because resolving 
+    // an imported name in a .py file in findSimilarElement() still prioritizes .pyi stubs over implementations
+    // as PyResolveContext doesn't retain the PyQualifiedNameResolveContext.getWithoutStubs flag.
+    // TODO propagate "without stubs" property through PyResolveContext
+    if (result instanceof PyElement && isInsideStub(result) && result.getContainingFile() != file) {
+      result = getOriginalElement((PyElement)result);
+    }
+
+    if (result != null) return result;
+
+    if (element instanceof PyFunction) {
      PyClass containingClass = PyUtil.turnConstructorIntoClass((PyFunction)element);
      if (containingClass != null) {
        result = findSimilarElement(containingClass, originalFile);
@@ -147,7 +160,7 @@ public final class PyiUtil {
  }

  public static boolean isPyiFileOfPackage(@NotNull PsiElement element) {
-    return element instanceof PyiFile || PyUtil.turnDirIntoInit(element) instanceof PyiFile;
+    return element instanceof PyiFile || PyUtil.turnDirIntoInitPyi(element) instanceof PyiFile;
  }

  private static boolean pyButNotPyiFile(@Nullable PsiFile file) {
@@ -164,7 +177,7 @@ public final class PyiUtil {
    return PyUtil.as(PyResolveImportUtil.resolveQualifiedName(name, context)
      .stream()
      .findFirst()
-      .map(PyUtil::turnDirIntoInit)
+      .map(PyUtil::turnDirIntoInitPyi)
      .orElse(null), PyiFile.class);
  }

@@ -178,7 +191,7 @@ public final class PyiUtil {
    return PyUtil.as(PyResolveImportUtil.resolveQualifiedName(name, context)
                       .stream()
                       .findFirst()
-                       .map(PyUtil::turnDirIntoInit)
+                       .map(PyUtil::turnDirIntoInitPy)
                       .orElse(null), PyFile.class);
  }

@@ -0,0 +1,2 @@
+def bar():
+    pass
@@ -0,0 +1 @@
+def bar() -> None: ...
@@ -0,0 +1,4 @@
+from pkg.foo import bar
+
+if __name__ == '__main__':
+    ba<caret>r()
@@ -0,0 +1 @@
+from blues.client import Blues
@@ -0,0 +1,3 @@
+from . import client
+
+Blues = client.Blues
@@ -0,0 +1,2 @@
+class Blues:
+    ...
@@ -0,0 +1,2 @@
+class Blues:
+    ...
@@ -0,0 +1,2 @@
+import blues
+blues.Blu<caret>es
@@ -0,0 +1 @@
+from .mod import Class
@@ -0,0 +1,3 @@
+from . import mod
+
+Class = mod.Class
@@ -0,0 +1 @@
+from .mod2 import Class
@@ -0,0 +1,3 @@
+from . import mod2
+
+Class = mod2.Class
@@ -0,0 +1,2 @@
+class Class:
+    pass
@@ -0,0 +1,2 @@
+class Class:
+    pass
@@ -0,0 +1,3 @@
+import pkg
+
+pkg.Cl<caret>ass
@@ -0,0 +1 @@
+import mymodule
@@ -0,0 +1,2 @@
+class Clazz:
+    ...
@@ -0,0 +1,2 @@
+class Clazz:
+    ...
@@ -0,0 +1,2 @@
+import pkg.mymodule
+pkg.mymodule.Claz<caret>z
@@ -0,0 +1 @@
+from . import mymodule
@@ -0,0 +1 @@
+from . import mymodule
@@ -0,0 +1,2 @@
+class Clazz:
+    ...
@@ -0,0 +1,2 @@
+class Clazz:
+    ...
@@ -0,0 +1,2 @@
+import pkg.mymodule
+pkg.mymodule.Claz<caret>z
@@ -0,0 +1,5 @@
+class Parent:
+    pass
+
+class Child(Parent):
+    pass
@@ -0,0 +1,6 @@
+class Parent:
+    ...
+
+
+class Child(Parent):
+    ...
@@ -0,0 +1,3 @@
+from mod import Parent
+
+print(Par<caret>ent)
@@ -0,0 +1,2 @@
+def bar():
+    pass
@@ -0,0 +1 @@
+def bar() -> None: ...
@@ -0,0 +1,4 @@
+from pkg.mymodule import bar
+
+if __name__ == '__main__':
+    ba<caret>r()
@@ -0,0 +1,3 @@
+from .mymodule import foo
+
+aaa = foo
@@ -0,0 +1,4 @@
+def foo():
+    ...
+
+aaa = foo
@@ -0,0 +1,2 @@
+def foo():
+    print(42)
@@ -0,0 +1,3 @@
+import pkg
+
+pkg.a<caret>aa
@@ -0,0 +1,8 @@
+class Parent:
+    def show(self):
+        pass
+
+
+class Child(Parent):
+    def show(self):
+        pass
@@ -0,0 +1,8 @@
+class Parent:
+    def show(self):
+        ...
+
+
+class Child(Parent):
+    def show(self):
+        ...
@@ -0,0 +1,5 @@
+from mod import Parent
+
+
+def foo(p: Parent):
+    p.sh<caret>ow()
@@ -0,0 +1 @@
+from blues.client import Blues
@@ -0,0 +1,3 @@
+from . import client
+
+Blues = client.Blues
@@ -0,0 +1,2 @@
+class Blues:
+    ...
@@ -0,0 +1,2 @@
+class Blues:
+    ...
@@ -0,0 +1,2 @@
+import blues
+blues.Blu<caret>es
@@ -0,0 +1 @@
+from .mod import Class
@@ -0,0 +1,3 @@
+from . import mod
+
+Class = mod.Class
@@ -0,0 +1 @@
+from .mod2 import Class
@@ -0,0 +1,3 @@
+from . import mod2
+
+Class = mod2.Class
@@ -0,0 +1,2 @@
+class Class:
+    pass
@@ -0,0 +1,2 @@
+class Class:
+    pass
@@ -0,0 +1,3 @@
+import pkg
+
+pkg.Cl<caret>ass
@@ -218,6 +218,73 @@ class PyNavigationTest : PyTestCase() {
    checkPyNotPyi(target?.containingFile)
  }

+  // PY-54905
+  fun testGoToImplementationFunctionInPackageWithInitPy() {
+    doTestGotoImplementationNavigatesToPyNotPyi()
+  }
+
+  // PY-54905
+  fun testGoToImplementationClassInPackageWithInitPy() {
+    doTestGotoImplementationNavigatesToPyNotPyi()
+  }
+
+  // PY-54905 PY-54620
+  fun testGoToImplementationClassInPackageWithInitPyi() {
+    doTestGotoImplementationNavigatesToPyNotPyi()
+  }
+
+  // PY-54905
+  fun testGoToImplementationFunctionInPyNotPyi() {
+    doTestGotoImplementationNavigatesToPyNotPyi(2)
+  }
+
+  // PY-54905
+  fun testGoToImplementationNameReExportedThroughAssignmentInPyiStub() {
+    doTestGotoImplementationNavigatesToPyNotPyi()
+  }
+
+  // PY-61740
+  fun testGoToDeclarationNameReExportedThroughAssignmentInPyiStub() {
+    doTestGotoDeclarationNavigatesToPyNotPyi()
+  }
+
+  // PY-61740
+  fun testGoToDeclarationNameReExportedThroughAssignmentInPyiStubTwice() {
+    doTestGotoDeclarationNavigatesToPyNotPyi()
+  }
+
+  // PY-54905
+  fun testGoToImplementationFunctionOverrides() {
+    doTestGotoImplementationNavigatesToPyNotPyi(2)
+  }
+
+  // PY-54905
+  fun testGoToImplementationClassInherits() {
+    doTestGotoImplementationNavigatesToPyNotPyi(2)
+  }
+
+  // PY-61740
+  fun testGoToDeclarationClassInPackageWithInitPyi() {
+    doTestGotoDeclarationNavigatesToPyNotPyi()
+  }
+
+  private fun doTestGotoDeclarationNavigatesToPyNotPyi() {
+    myFixture.copyDirectoryToProject(getTestName(true), "")
+    myFixture.configureByFile("test.py")
+    val target = PyGotoDeclarationHandler().getGotoDeclarationTarget(elementAtCaret, myFixture.editor)
+    checkPyNotPyi(target!!.containingFile)
+  }
+
+  private fun doTestGotoImplementationNavigatesToPyNotPyi(numTargets: Int = 1) {
+    myFixture.copyDirectoryToProject(getTestName(true), "")
+    myFixture.configureByFile("test.py")
+    val gotoData = CodeInsightTestUtil.gotoImplementation(myFixture.editor, myFixture.file)
+    assertSize(numTargets, gotoData.targets)
+    for (target in gotoData.targets) {
+      checkPyNotPyi(target.containingFile)
+    }
+  }
+
  private fun doTestGotoDeclarationOrUsagesOutcome(expectedOutcome: GTDUOutcome, text: String) {
    myFixture.configureByText("a.py", text)
    val actualOutcome = GotoDeclarationOrUsageHandler2.testGTDUOutcomeInNonBlockingReadAction(myFixture.editor, myFixture.file, myFixture.caretOffset)