From 7ce9158c269bbedbec098edbe351a74c1bf93a22 Mon Sep 17 00:00:00 2001 From: Morgan Bartholomew Date: Tue, 17 Jun 2025 21:45:45 +1000 Subject: [PATCH] [python] refactor stub suggestions to use `PythonPackageManager` (cherry picked from commit d370e9fb005f850246a377ae10e4e87ba924efbd) GitOrigin-RevId: aea57ce2819119f7f8cb877def6f13a65dc5e754 --- .../python/packaging/PyPackageManagers.java | 3 +- .../typing/PyStubPackagesAdvertiser.kt | 81 ++++++++++--------- .../typing/PyStubPackagesAdvertiserCache.kt | 9 ++- .../PyStubPackagesCompatibilityInspection.kt | 3 +- .../typing/PyStubPackagesLoader.kt | 19 ++--- 5 files changed, 62 insertions(+), 53 deletions(-) diff --git a/python/openapi/src/com/jetbrains/python/packaging/PyPackageManagers.java b/python/openapi/src/com/jetbrains/python/packaging/PyPackageManagers.java index 5f23a426ccf0..f104e95fc6ec 100644 --- a/python/openapi/src/com/jetbrains/python/packaging/PyPackageManagers.java +++ b/python/openapi/src/com/jetbrains/python/packaging/PyPackageManagers.java @@ -9,8 +9,7 @@ import com.intellij.webcore.packaging.PackageManagementService; import org.jetbrains.annotations.NotNull; /** - * @deprecated replaced by {@link com.jetbrains.python.packaging.common.PackageManagerHolder }. - * To get an instance of PythonPackageManager consider using + * @deprecated To get an instance of PythonPackageManager consider using * {@link com.jetbrains.python.packaging.management.PythonPackageManager.Companion#forSdk(Project, Sdk)} */ @Deprecated(forRemoval = true) diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt index 18392f3a377a..f2137af00077 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt @@ -13,7 +13,7 @@ import com.intellij.notification.NotificationAction import com.intellij.notification.NotificationGroupManager import com.intellij.notification.NotificationType import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.components.ComponentManager +import com.intellij.openapi.components.service import com.intellij.openapi.module.Module import com.intellij.openapi.module.ModuleUtilCore import com.intellij.openapi.project.Project @@ -30,6 +30,7 @@ import com.jetbrains.python.inspections.PyInspectionVisitor import com.jetbrains.python.inspections.quickfix.PyInstallRequirementsFix import com.jetbrains.python.inspections.requirement.RunningPackagingTasksListener import com.jetbrains.python.packaging.* +import com.jetbrains.python.packaging.management.PythonPackageManager import com.jetbrains.python.packaging.requirement.PyRequirementRelation import com.jetbrains.python.psi.PyFile import com.jetbrains.python.psi.PyReferenceExpression @@ -104,19 +105,19 @@ private class PyStubPackagesAdvertiser : PyInspection() { val module = ModuleUtilCore.findModuleForFile(file) ?: return val sdk = PythonSdkUtil.findPythonSdk(module) ?: return - val packageManager = PyPackageManager.getInstance(sdk) - val installedPackages = packageManager.packages ?: emptyList() + val packageManager = PythonPackageManager.forSdk(module.project, sdk) + val installedPackages = packageManager.listInstalledPackagesSnapshot() if (installedPackages.isEmpty()) return val packageManagementService = PyPackageManagers.getInstance().getManagementService(file.project, sdk) val availablePackages = packageManagementService.allPackagesCached if (availablePackages.isEmpty()) return - val ignoredStubPackages = (IGNORE + ignoredPackages).mapNotNull { packageManager.parseRequirement(it) } - val cache = ApplicationManager.getApplication().getService().forSdk(sdk) + val ignoredStubPackages = (IGNORE + ignoredPackages).mapNotNull { PyRequirementParser.fromLine(it) } + val cache = ApplicationManager.getApplication().service().forSdk(sdk) - val forcedToLoad = processForcedPackages(file, sources, module, sdk, packageManager, ignoredStubPackages, cache) - val checkedToLoad = processCheckedPackages(file, sources, module, sdk, packageManager, ignoredStubPackages, cache) + val forcedToLoad = processForcedPackages(file, sources, module, sdk, ignoredStubPackages, cache) + val checkedToLoad = processCheckedPackages(file, sources, module, sdk, ignoredStubPackages, cache) loadStubPackagesForSources( forcedToLoad + checkedToLoad, @@ -128,13 +129,14 @@ private class PyStubPackagesAdvertiser : PyInspection() { ) } - private fun processForcedPackages(file: PyFile, - sources: Set, - module: Module, - sdk: Sdk, - packageManager: PyPackageManager, - ignoredStubPackages: List, - cache: Cache): Set { + private fun processForcedPackages( + file: PyFile, + sources: Set, + module: Module, + sdk: Sdk, + ignoredStubPackages: List, + cache: Cache + ): Set { val (sourcesToLoad, cached) = splitIntoNotCachedAndCached(forcedSourcesToProcess(sources), cache) val (reqs, args) = toRequirementsAndExtraArgs(cached, ignoredStubPackages) @@ -144,20 +146,21 @@ private class PyStubPackagesAdvertiser : PyInspection() { registerProblem(file, message, - createInstallStubPackagesQuickFix(reqs, args, module, sdk, packageManager), - createIgnorePackagesQuickFix(reqs, packageManager)) + createInstallStubPackagesQuickFix(reqs, args, module, sdk), + createIgnorePackagesQuickFix(reqs)) } return sourcesToLoad } - private fun processCheckedPackages(file: PyFile, - sources: Set, - module: Module, - sdk: Sdk, - packageManager: PyPackageManager, - ignoredStubPackages: List, - cache: Cache): Set { + private fun processCheckedPackages( + file: PyFile, + sources: Set, + module: Module, + sdk: Sdk, + ignoredStubPackages: List, + cache: Cache + ): Set { val project = file.project if (project.getUserData(BALLOON_SHOWING) == true) return emptySet() @@ -165,7 +168,7 @@ private class PyStubPackagesAdvertiser : PyInspection() { val (unfilteredReqs, args) = toRequirementsAndExtraArgs(cached, ignoredStubPackages) - val status = file.project.getService() + val status = file.project.service() val reqs = unfilteredReqs.filterNot { status.markedAsInstalling(it.name) } @@ -197,11 +200,11 @@ private class PyStubPackagesAdvertiser : PyInspection() { NotificationAction.createSimpleExpiring( if (plural) PyBundle.message("code.insight.install.type.hints.action") else "${PyBundle.message("python.packaging.install")} $reqsToString" - ) { createInstallStubPackagesQuickFix(reqs, args, module, sdk, packageManager).applyFix(project, problemDescriptor) } + ) { createInstallStubPackagesQuickFix(reqs, args, module, sdk).applyFix(project, problemDescriptor) } ) .addAction( NotificationAction.createSimpleExpiring(PyBundle.message("code.insight.ignore.type.hints")) { - createIgnorePackagesQuickFix(reqs, packageManager).applyFix(project, problemDescriptor) + createIgnorePackagesQuickFix(reqs).applyFix(project, problemDescriptor) } ) .addAction( @@ -271,21 +274,22 @@ private class PyStubPackagesAdvertiser : PyInspection() { return requirements to args.toList() } - private fun createInstallStubPackagesQuickFix(reqs: List, - args: List, - module: Module, - sdk: Sdk, - packageManager: PyPackageManager): LocalQuickFix { + private fun createInstallStubPackagesQuickFix( + reqs: List, + args: List, + module: Module, + sdk: Sdk, + ): LocalQuickFix { val project = module.project val stubPkgNamesToInstall = reqs.mapTo(mutableSetOf()) { it.name } val installationListener = object : RunningPackagingTasksListener(module) { override fun started() { - project.getService().markAsInstalling(stubPkgNamesToInstall) + project.service().markAsInstalling(stubPkgNamesToInstall) } override fun finished(exceptions: List) { - val status = project.getService() + val status = project.service() val stubPkgsToUninstall = PyStubPackagesCompatibilityInspection .findIncompatibleRuntimeToStubPackages(sdk) { it.name in stubPkgNamesToInstall } @@ -300,7 +304,7 @@ private class PyStubPackagesAdvertiser : PyInspection() { override fun finished(exceptions: MutableList?) { status.unmarkAsInstalling(stubPkgNamesToUninstall) val reqsToIgnore = stubPkgsToUninstall.map { pyRequirement(it.name, PyRequirementRelation.EQ, it.version) } - addStubPackagesToIgnore(reqsToIgnore, stubPkgNamesToUninstall, project, packageManager) + addStubPackagesToIgnore(reqsToIgnore, stubPkgNamesToUninstall, project) } } @@ -321,21 +325,20 @@ private class PyStubPackagesAdvertiser : PyInspection() { return PyInstallRequirementsFix(name, sdk, reqs, args, installationListener) } - private fun createIgnorePackagesQuickFix(reqs: List, packageManager: PyPackageManager): LocalQuickFix { + private fun createIgnorePackagesQuickFix(reqs: List): LocalQuickFix { return object : LocalQuickFix { override fun getFamilyName() = PyBundle.message("code.insight.ignore.packages.qfix", reqs.size) override fun applyFix(project: Project, descriptor: ProblemDescriptor) { - this@Visitor.addStubPackagesToIgnore(reqs, reqs.mapTo(mutableSetOf()) { it.name }, project, packageManager) + this@Visitor.addStubPackagesToIgnore(reqs, reqs.mapTo(mutableSetOf()) { it.name }, project) } } } private fun addStubPackagesToIgnore(stubPackages: List, stubPackagesNames: Set, - project: Project, - packageManager: PyPackageManager) { - ignoredPackages.removeIf { packageManager.parseRequirement(it)?.name in stubPackagesNames } + project: Project) { + ignoredPackages.removeIf { PyRequirementParser.fromLine(it)?.name in stubPackagesNames } ignoredPackages.addAll(stubPackages.map { it.presentableText }) ProjectInspectionProfileManager.getInstance(project).fireProfileChanged() diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiserCache.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiserCache.kt index 0591f806045f..0b64dab72b29 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiserCache.kt +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiserCache.kt @@ -8,7 +8,8 @@ import com.google.common.cache.LoadingCache import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.components.Service import com.intellij.openapi.projectRoots.Sdk -import com.jetbrains.python.packaging.PyPackageManager +import com.jetbrains.python.packaging.common.PythonPackageManagementListener +import com.jetbrains.python.packaging.management.PythonPackageManager.Companion.PACKAGE_MANAGEMENT_TOPIC import java.time.Duration @Service @@ -28,7 +29,11 @@ class PyStubPackagesAdvertiserCache { init { val connection = ApplicationManager.getApplication().messageBus.connect() - connection.subscribe(PyPackageManager.PACKAGE_MANAGER_TOPIC, PyPackageManager.Listener { cache.invalidate(it) }) + connection.subscribe(PACKAGE_MANAGEMENT_TOPIC, object : PythonPackageManagementListener { + override fun packagesChanged(sdk: Sdk) { + cache.invalidate(sdk) + } + }) } fun forSdk(sdk: Sdk): Cache { diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesCompatibilityInspection.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesCompatibilityInspection.kt index af75d98d1684..d58eda17db50 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesCompatibilityInspection.kt +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesCompatibilityInspection.kt @@ -6,6 +6,7 @@ import com.intellij.codeInspection.LocalQuickFix import com.intellij.codeInspection.ProblemDescriptor import com.intellij.codeInspection.ProblemsHolder import com.intellij.codeInspection.options.OptPane +import com.intellij.openapi.components.service import com.intellij.openapi.module.ModuleUtilCore import com.intellij.openapi.project.Project import com.intellij.openapi.projectRoots.Sdk @@ -77,7 +78,7 @@ class PyStubPackagesCompatibilityInspection : PyInspection() { val nameToPkg = mutableMapOf() installedPackages.forEach { nameToPkg[it.name] = it } - val status = node.project.getService(PyStubPackagesInstallingStatus::class.java) + val status = node.project.service() findIncompatibleRuntimeToStubPackages( sdk) { stubPkg -> diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesLoader.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesLoader.kt index 80592f640c79..f89bd58380a3 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesLoader.kt +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesLoader.kt @@ -2,19 +2,20 @@ package com.jetbrains.python.codeInsight.typing import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.components.service import com.intellij.openapi.projectRoots.Sdk import com.intellij.util.CatchingConsumer import com.intellij.webcore.packaging.PackageManagementService import com.intellij.webcore.packaging.RepoPackage import com.jetbrains.python.packaging.PyPIPackageUtil -import com.jetbrains.python.packaging.PyPackage +import com.jetbrains.python.packaging.common.PythonPackage import java.util.* import java.util.function.BiConsumer fun loadStubPackagesForSources(sourcesToLoad: Set, sourceToPackage: Map, - installedPackages: List, + installedPackages: List, availablePackages: List, packageManagementService: PackageManagementService, sdk: Sdk) { @@ -27,15 +28,15 @@ fun loadStubPackagesForSources(sourcesToLoad: Set, sourceToStubPackagesAvailableToInstall, packageManagementService, BiConsumer { source, stubPackagesForSource -> - ApplicationManager.getApplication().getService(PyStubPackagesAdvertiserCache::class.java).forSdk(sdk).put(source, stubPackagesForSource) + ApplicationManager.getApplication().service().forSdk(sdk).put(source, stubPackagesForSource) } ) } private fun sourceToInstalledRuntimeAndStubPackages(sourcesToLoad: Set, sourceToPackage: Map, - installedPackages: List): Map>> { - val result = mutableMapOf>>() + installedPackages: List): Map>> { + val result = mutableMapOf>>() for (source in sourcesToLoad) { val pkgName = sourceToPackage[source] ?: continue @@ -45,7 +46,7 @@ private fun sourceToInstalledRuntimeAndStubPackages(sourcesToLoad: Set, return result } -private fun sourceToStubPackagesAvailableToInstall(sourceToInstalledRuntimeAndStubPkgs: Map>>, +private fun sourceToStubPackagesAvailableToInstall(sourceToInstalledRuntimeAndStubPkgs: Map>>, availablePackages: List): Map> { if (sourceToInstalledRuntimeAndStubPkgs.isEmpty()) return emptyMap() @@ -83,9 +84,9 @@ private fun loadRequirementsAndExtraArgs(sourceToStubPackagesAvailableToInstall: } } -private fun installedRuntimeAndStubPackages(pkgName: String, installedPackages: List): Pair? { - var runtime: PyPackage? = null - var stub: PyPackage? = null +private fun installedRuntimeAndStubPackages(pkgName: String, installedPackages: List): Pair? { + var runtime: PythonPackage? = null + var stub: PythonPackage? = null val stubPkgName = "$pkgName$STUBS_SUFFIX" val typesPkgName = "$TYPES_PREFIX$pkgName" val typesSuffixPkgName = "$pkgName$TYPES_SUFFIX"