diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyBundledStubs.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyBundledStubs.kt new file mode 100644 index 000000000000..58a2c1761c19 --- /dev/null +++ b/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 { + return listOfNotNull(BUNDLED_STUB_ROOT); + } +} diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt b/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt index e2c054494e74..f7ba1939cb0f 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt +++ b/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 diff --git a/python/src/com/jetbrains/python/sdk/PythonSdkUpdater.java b/python/src/com/jetbrains/python/sdk/PythonSdkUpdater.java index 68c824b5e405..a7bfd3b8783e 100644 --- a/python/src/com/jetbrains/python/sdk/PythonSdkUpdater.java +++ b/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 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. */