PY-76629: Do not resolve bundled stubs when the package is not installed

The implementation is very similar to the PyTypeShed,
it should be merged into more general solution.

GitOrigin-RevId: 1449a774e93206349f17f264c27e2a7877d79b19

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyBundledStubs.kt

@@ -0,0 +1,57 @@
+package com.jetbrains.python.codeInsight.typing
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.vfs.StandardFileSystems
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.util.QualifiedName
+import com.intellij.util.PlatformUtils
+import com.jetbrains.python.PyPsiPackageUtil
+import com.jetbrains.python.PythonHelpersLocator
+import com.jetbrains.python.packaging.PyPackageManagers
+
+object PyBundledStubs {
+  private const val BUNDLED_STUBS_PATH = "bundled_stubs"
+
+  /**
+   * The actual bundled stubs directory.
+   */
+  private val BUNDLED_STUB_ROOT: VirtualFile? by lazy {
+    var helpersPath = PythonHelpersLocator.findPathStringInHelpers(BUNDLED_STUBS_PATH);
+    if (helpersPath.isEmpty()) {
+      return@lazy null
+    }
+    StandardFileSystems.local().findFileByPath(helpersPath)
+  }
+
+  fun maySearchForStubInRoot(name: QualifiedName, root: VirtualFile, sdk: Sdk): Boolean {
+    // TODO merge with PyTypeShed
+    if (ApplicationManager.getApplication().isUnitTestMode) {
+      return true
+    }
+    val possiblePackage = name.firstComponent ?: return false
+    val alternativePossiblePackage = PyPsiPackageUtil.moduleToPackageName(possiblePackage, default = "")
+
+    val packageManager = PyPackageManagers.getInstance().forSdk(sdk)
+    val installedPackages = if (ApplicationManager.getApplication().isHeadlessEnvironment && !PlatformUtils.isFleetBackend()) {
+      packageManager.refreshAndGetPackages(false)
+    }
+    else {
+      packageManager.packages ?: return true
+    }
+
+    return packageManager.parseRequirement(possiblePackage)?.match(installedPackages) != null ||
+           PyPsiPackageUtil.findPackage(installedPackages, alternativePossiblePackage) != null
+  }
+
+  /**
+   * Checks if the [file] is located inside the bundled stubs directory.
+   */
+  fun isBundledStubsDirectory(file: VirtualFile): Boolean {
+    return file == BUNDLED_STUB_ROOT
+  }
+
+  fun getRoots(): Iterable<VirtualFile> {
+    return listOfNotNull(BUNDLED_STUB_ROOT);
+  }
+}

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

@@ -22,6 +22,7 @@ import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.QualifiedName
 import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.intellij.util.concurrency.annotations.RequiresReadLock
+import com.jetbrains.python.codeInsight.typing.PyBundledStubs
 import com.jetbrains.python.codeInsight.typing.PyTypeShed
 import com.jetbrains.python.codeInsight.typing.isInInlinePackage
 import com.jetbrains.python.codeInsight.typing.isInStubPackage
@@ -280,6 +281,9 @@ private fun resultsFromRoots(name: QualifiedName, context: PyQualifiedNameResolv
         effectiveSdk != null && PyTypeShed.isInside(root) && !PyTypeShed.maySearchForStubInRoot(name, root, effectiveSdk)) {
       return@RootVisitor true
     }
+    if (effectiveSdk != null && PyBundledStubs.isBundledStubsDirectory(root) && !PyBundledStubs.maySearchForStubInRoot(name, root, effectiveSdk)) {
+      return@RootVisitor true
+    }
     if (withoutStubs && (PyTypeShed.isInside(root) ||
                          PsiManager.getInstance(context.project).findDirectory(root)?.let { isInStubPackage(it) } == true)) {
       return@RootVisitor true

python/src/com/jetbrains/python/sdk/PythonSdkUpdater.java

@@ -38,8 +38,8 @@ import com.intellij.util.ExceptionUtil;
 import com.intellij.util.PathMappingSettings;
 import com.intellij.util.Processor;
 import com.jetbrains.python.PyBundle;
-import com.jetbrains.python.PythonHelpersLocator;
 import com.jetbrains.python.PythonPluginDisposable;
+import com.jetbrains.python.codeInsight.typing.PyBundledStubs;
 import com.jetbrains.python.codeInsight.typing.PyTypeShed;
 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
 import com.jetbrains.python.packaging.PyPackageManager;
@@ -503,20 +503,10 @@ public final class PythonSdkUpdater {
       .addAll(getSkeletonsPaths(sdk))
       .addAll(userAddedRoots)
       .addAll(PyTypeShed.INSTANCE.findRootsForSdk(sdk))
-      .addAll(getBundledStubs())
+      .addAll(PyBundledStubs.INSTANCE.getRoots())
       .build();
   }
 
-  private static @NotNull List<VirtualFile> getBundledStubs() {
-    var helpersPath = PythonHelpersLocator.findPathStringInHelpers("bundled_stubs");
-    if (helpersPath.isEmpty()) {
-      return Collections.emptyList();
-    }
-    VirtualFile bundledStubRoot = StandardFileSystems.local().findFileByPath(helpersPath);
-    if (bundledStubRoot == null) return Collections.emptyList();
-    return List.of(bundledStubRoot);
-  }
-
   /**
    * Returns all the paths manually added to an SDK by the user.
    */