diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt index 0910a8d1d50e..1e7b04f69605 100644 --- a/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt +++ b/python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt @@ -19,7 +19,6 @@ import com.intellij.openapi.util.Key import com.intellij.profile.codeInspection.ProjectInspectionProfileManager import com.intellij.psi.PsiElementVisitor import com.intellij.psi.util.QualifiedName -import com.intellij.util.containers.isNullOrEmpty import com.jetbrains.python.codeInsight.typing.PyStubPackagesAdvertiserCache.Companion.StubPackagesForSource import com.jetbrains.python.inspections.PyInspection import com.jetbrains.python.inspections.PyInspectionVisitor @@ -27,8 +26,6 @@ import com.jetbrains.python.inspections.PyPackageRequirementsInspection.PyInstal import com.jetbrains.python.packaging.* import com.jetbrains.python.packaging.requirement.PyRequirementRelation import com.jetbrains.python.psi.PyFile -import com.jetbrains.python.psi.PyFromImportStatement -import com.jetbrains.python.psi.PyImportElement import com.jetbrains.python.psi.PyReferenceExpression import com.jetbrains.python.sdk.PythonSdkUtil import javax.swing.JComponent @@ -49,8 +46,6 @@ class PyStubPackagesAdvertiser : PyInspection() { private val BALLOON_SHOWING = Key.create("showingStubPackagesAdvertiserBalloon") private val BALLOON_NOTIFICATIONS = NotificationGroup("Python Stub Packages Advertiser", NotificationDisplayType.STICKY_BALLOON, true) - - private val SESSION_KEY = Key.create>("PyStubPackagesAdvertiser.Sources") } @Suppress("MemberVisibilityCanBePrivate") @@ -60,287 +55,278 @@ class PyStubPackagesAdvertiser : PyInspection() { override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean, - session: LocalInspectionToolSession): PsiElementVisitor { - val sources = session.putUserDataIfAbsent(SESSION_KEY, mutableSetOf()) - return Visitor(sources, holder, session) - } + session: LocalInspectionToolSession): PsiElementVisitor = Visitor(ignoredPackages, holder, session) - override fun inspectionFinished(session: LocalInspectionToolSession, problemsHolder: ProblemsHolder) { - val sources = session.getUserData(SESSION_KEY) - if (sources.isNullOrEmpty()) return - - val file = session.file - if (file is PyFile) run(file, sources!!, problemsHolder) - } - - private class Visitor(val sources: MutableSet, + private class Visitor(private val ignoredPackages: MutableList, holder: ProblemsHolder, session: LocalInspectionToolSession) : PyInspectionVisitor(holder, session) { - override fun visitPyFromImportStatement(node: PyFromImportStatement) { - super.visitPyFromImportStatement(node) - processImport(node.importSource, node.importSourceQName) + override fun visitPyFile(node: PyFile) { + super.visitPyFile(node) + + val sources = mutableSetOf() + + node.fromImports.mapNotNullTo(sources) { topLevelPackagesWithoutStubs(it.importSource, it.importSourceQName) } + node.importTargets.mapNotNullTo(sources) { topLevelPackagesWithoutStubs(it.importReferenceExpression, it.importedQName) } + + if (sources.isNotEmpty()) { + run(node, sources) + } } - override fun visitPyImportElement(node: PyImportElement) { - super.visitPyImportElement(node) - processImport(node.importReferenceExpression, node.importedQName) - } - - private fun processImport(ref: PyReferenceExpression?, qName: QualifiedName?) { - if (qName == null) return + private fun topLevelPackagesWithoutStubs(ref: PyReferenceExpression?, qName: QualifiedName?): String? { + if (qName == null) return null if (ref != null && ref.getReference(resolveContext).multiResolve(false).asSequence().mapNotNull { it.element }.any { isInStubPackage(it) }) { - return + return null } - qName.firstComponent?.let(sources::add) - } - } - - private fun run(file: PyFile, sources: Set, problemsHolder: ProblemsHolder) { - val module = ModuleUtilCore.findModuleForFile(file) ?: return - val sdk = PythonSdkUtil.findPythonSdk(module) ?: return - - val packageManager = PyPackageManager.getInstance(sdk) - val installedPackages = packageManager.packages ?: emptyList() - if (installedPackages.isEmpty()) return - - val packageManagementService = PyPackageManagers.getInstance().getManagementService(file.project, sdk) - val availablePackages = packageManagementService.allPackagesCached - if (availablePackages.isEmpty()) return - - val ignoredStubPackages = ignoredPackages.mapNotNull { packageManager.parseRequirement(it) } - val cache = ServiceManager.getService(PyStubPackagesAdvertiserCache::class.java).forSdk(sdk) - - val forcedToLoad = processForcedPackages(file, sources, module, sdk, packageManager, ignoredStubPackages, cache, problemsHolder) - val checkedToLoad = processCheckedPackages(file, sources, module, sdk, packageManager, ignoredStubPackages, cache) - - loadStubPackagesForSources( - forcedToLoad + checkedToLoad, - FORCED + CHECKED, - installedPackages, - availablePackages, - packageManagementService, - sdk - ) - } - - private fun processForcedPackages(file: PyFile, - sources: Set, - module: Module, - sdk: Sdk, - packageManager: PyPackageManager, - ignoredStubPackages: List, - cache: Cache, - problemsHolder: ProblemsHolder): Set { - val (sourcesToLoad, cached) = splitIntoNotCachedAndCached(forcedSourcesToProcess(sources), cache) - - val (reqs, args) = toRequirementsAndExtraArgs(cached, ignoredStubPackages) - if (reqs.isNotEmpty()) { - val plural = reqs.size > 1 - val reqsToString = PyPackageUtil.requirementsToString(reqs) - - problemsHolder.registerProblem(file, - "Stub package${if (plural) "s" else ""} $reqsToString ${if (plural) "are" else "is"} not installed. " + - "${if (plural) "They" else "It"} contain${if (plural) "" else "s"} type hints needed for better code insight.", - createInstallStubPackagesQuickFix(reqs, args, module, sdk, packageManager), - createIgnorePackagesQuickFix(reqs, packageManager)) + return qName.firstComponent } - return sourcesToLoad - } + private fun run(file: PyFile, sources: Set) { + val module = ModuleUtilCore.findModuleForFile(file) ?: return + val sdk = PythonSdkUtil.findPythonSdk(module) ?: return - private fun processCheckedPackages(file: PyFile, - sources: Set, - module: Module, - sdk: Sdk, - packageManager: PyPackageManager, - ignoredStubPackages: List, - cache: Cache): Set { - val project = file.project - if (project.getUserData(BALLOON_SHOWING) == true) return emptySet() + val packageManager = PyPackageManager.getInstance(sdk) + val installedPackages = packageManager.packages ?: emptyList() + if (installedPackages.isEmpty()) return - val (sourcesToLoad, cached) = splitIntoNotCachedAndCached(checkedSourcesToProcess(sources), cache) + val packageManagementService = PyPackageManagers.getInstance().getManagementService(file.project, sdk) + val availablePackages = packageManagementService.allPackagesCached + if (availablePackages.isEmpty()) return - val (reqs, args) = toRequirementsAndExtraArgs(cached, ignoredStubPackages) - if (reqs.isNotEmpty()) { - val plural = reqs.size > 1 - val reqsToString = PyPackageUtil.requirementsToString(reqs) + val ignoredStubPackages = ignoredPackages.mapNotNull { packageManager.parseRequirement(it) } + val cache = ServiceManager.getService(PyStubPackagesAdvertiserCache::class.java).forSdk(sdk) - project.putUserData(BALLOON_SHOWING, true) + val forcedToLoad = processForcedPackages(file, sources, module, sdk, packageManager, ignoredStubPackages, cache) + val checkedToLoad = processCheckedPackages(file, sources, module, sdk, packageManager, ignoredStubPackages, cache) - BALLOON_NOTIFICATIONS - .createNotification( - "Type hints are not installed", - "They could make code insight better.
" + - "Install ${if (plural) "stub packages" else reqsToString}    " + - "Ignore    " + - "Settings", - NotificationType.INFORMATION - ) { notification, event -> - try { - val problemDescriptor = ProblemDescriptorImpl( - file, - file, - "Stub package${if (plural) "s" else ""} $reqsToString ${if (plural) "are" else "is"} not installed", - LocalQuickFix.EMPTY_ARRAY, - ProblemHighlightType.GENERIC_ERROR_OR_WARNING, - true, - null, - true - ) + loadStubPackagesForSources( + forcedToLoad + checkedToLoad, + FORCED + CHECKED, + installedPackages, + availablePackages, + packageManagementService, + sdk + ) + } - when (event.description) { - "#yes" -> { - createInstallStubPackagesQuickFix(reqs, args, module, sdk, packageManager).applyFix(project, problemDescriptor) - } - "#no" -> createIgnorePackagesQuickFix(reqs, packageManager).applyFix(project, problemDescriptor) - "#settings" -> { - val profile = ProjectInspectionProfileManager.getInstance(project).currentProfile - EditInspectionToolsSettingsAction.editToolSettings(project, profile, PyStubPackagesAdvertiser::class.simpleName) + private fun processForcedPackages(file: PyFile, + sources: Set, + module: Module, + sdk: Sdk, + packageManager: PyPackageManager, + ignoredStubPackages: List, + cache: Cache): Set { + val (sourcesToLoad, cached) = splitIntoNotCachedAndCached(forcedSourcesToProcess(sources), cache) + + val (reqs, args) = toRequirementsAndExtraArgs(cached, ignoredStubPackages) + if (reqs.isNotEmpty()) { + val plural = reqs.size > 1 + val reqsToString = PyPackageUtil.requirementsToString(reqs) + + registerProblem(file, + "Stub package${if (plural) "s" else ""} $reqsToString ${if (plural) "are" else "is"} not installed. " + + "${if (plural) "They" else "It"} contain${if (plural) "" else "s"} type hints needed for better code insight.", + createInstallStubPackagesQuickFix(reqs, args, module, sdk, packageManager), + createIgnorePackagesQuickFix(reqs, packageManager)) + } + + return sourcesToLoad + } + + private fun processCheckedPackages(file: PyFile, + sources: Set, + module: Module, + sdk: Sdk, + packageManager: PyPackageManager, + ignoredStubPackages: List, + cache: Cache): Set { + val project = file.project + if (project.getUserData(BALLOON_SHOWING) == true) return emptySet() + + val (sourcesToLoad, cached) = splitIntoNotCachedAndCached(checkedSourcesToProcess(sources), cache) + + val (reqs, args) = toRequirementsAndExtraArgs(cached, ignoredStubPackages) + if (reqs.isNotEmpty()) { + val plural = reqs.size > 1 + val reqsToString = PyPackageUtil.requirementsToString(reqs) + + project.putUserData(BALLOON_SHOWING, true) + + BALLOON_NOTIFICATIONS + .createNotification( + "Type hints are not installed", + "They could make code insight better.
" + + "Install ${if (plural) "stub packages" else reqsToString}    " + + "Ignore    " + + "Settings", + NotificationType.INFORMATION + ) { notification, event -> + try { + val problemDescriptor = ProblemDescriptorImpl( + file, + file, + "Stub package${if (plural) "s" else ""} $reqsToString ${if (plural) "are" else "is"} not installed", + LocalQuickFix.EMPTY_ARRAY, + ProblemHighlightType.GENERIC_ERROR_OR_WARNING, + true, + null, + true + ) + + when (event.description) { + "#yes" -> { + createInstallStubPackagesQuickFix(reqs, args, module, sdk, packageManager).applyFix(project, problemDescriptor) + } + "#no" -> createIgnorePackagesQuickFix(reqs, packageManager).applyFix(project, problemDescriptor) + "#settings" -> { + val profile = ProjectInspectionProfileManager.getInstance(project).currentProfile + EditInspectionToolsSettingsAction.editToolSettings(project, profile, PyStubPackagesAdvertiser::class.simpleName) + } } } - } - finally { - notification.expire() - } - } - .whenExpired { project.putUserData(BALLOON_SHOWING, false) } - .notify(project) - } - - return sourcesToLoad - } - - private fun forcedSourcesToProcess(sources: Set) = sources.filterTo(mutableSetOf()) { it in FORCED } - - private fun checkedSourcesToProcess(sources: Set) = sources.filterTo(mutableSetOf()) { it in CHECKED } - - private fun splitIntoNotCachedAndCached(sources: Set, - cache: Cache): Pair, List> { - if (sources.isEmpty()) return emptySet() to emptyList() - - val notCached = mutableSetOf() - val cached = mutableListOf() - - synchronized(cache) { - // despite cache is thread-safe, - // here we have sync block to guarantee only one reader - // and as a result not run processing for sources that are already evaluating - - sources.forEach { source -> - cache.getIfPresent(source).let { - if (it == null) { - notCached.add(source) - - // mark this source as evaluating - // if source processing failed, this value would mean that such source was handled - cache.put(source, StubPackagesForSource.EMPTY) - } - else { - cached.add(it) - } - } - } - } - - return notCached to cached - } - - private fun toRequirementsAndExtraArgs(cached: List, - ignoredStubPackages: List): Pair, List> { - if (cached.isEmpty()) return emptyList() to emptyList() - - val requirements = cached.asSequence() - .flatMap { it.packages.entries.asSequence() } - .filterNot { isIgnoredStubPackage(it.key, it.value.first, ignoredStubPackages) } - .map { - pyRequirement(it.key, PyRequirementRelation.EQ, it.value.first) - } - .toList() - if (requirements.isEmpty()) return emptyList() to emptyList() - - val args = sequenceOf("--no-deps") + - cached.asSequence().flatMap { pkgs -> pkgs.packages.values.asSequence().map { it.second }.flatten() } - return requirements to args.toList() - } - - private fun createInstallStubPackagesQuickFix(reqs: List, - args: List, - module: Module, - sdk: Sdk, - packageManager: PyPackageManager): LocalQuickFix { - val project = module.project - val stubPkgNamesToInstall = reqs.mapTo(mutableSetOf()) { it.name } - - val installationListener = object : PyPackageManagerUI.Listener { - override fun started() { - ServiceManager.getService(project, PyStubPackagesInstallingStatus::class.java).markAsInstalling(stubPkgNamesToInstall) - } - - override fun finished(exceptions: MutableList?) { - val status = ServiceManager.getService(project, PyStubPackagesInstallingStatus::class.java) - - val stubPkgsToUninstall = PyStubPackagesCompatibilityInspection - .findIncompatibleRuntimeToStubPackages(sdk) { it.name in stubPkgNamesToInstall } - .map { it.second } - - if (stubPkgsToUninstall.isNotEmpty()) { - val stubPkgNamesToUninstall = stubPkgsToUninstall.mapTo(mutableSetOf()) { it.name } - - val uninstallationListener = object : PyPackageManagerUI.Listener { - override fun started() {} - - override fun finished(exceptions: MutableList?) { - status.unmarkAsInstalling(stubPkgNamesToUninstall) - val reqsToIgnore = stubPkgsToUninstall.map { pyRequirement(it.name, PyRequirementRelation.EQ, it.version) } - addStubPackagesToIgnore(reqsToIgnore, stubPkgNamesToUninstall, project, packageManager) + finally { + notification.expire() } } + .whenExpired { project.putUserData(BALLOON_SHOWING, false) } + .notify(project) + } - val plural = stubPkgNamesToUninstall.size > 1 - val content = "Suggested ${stubPkgNamesToUninstall.joinToString { "'$it'" }} " + - "${if (plural) "are" else "is"} incompatible with your current environment.
" + - "${if (plural) "These" else "This"} stub package${if (plural) "s" else ""} will be removed and ignored until new version is released." + return sourcesToLoad + } - BALLOON_NOTIFICATIONS.createNotification(content, NotificationType.WARNING).notify(project) - PyPackageManagerUI(project, sdk, uninstallationListener).uninstall(stubPkgsToUninstall) + private fun forcedSourcesToProcess(sources: Set) = sources.filterTo(mutableSetOf()) { it in FORCED } - stubPkgNamesToInstall.removeAll(stubPkgNamesToUninstall) + private fun checkedSourcesToProcess(sources: Set) = sources.filterTo(mutableSetOf()) { it in CHECKED } + + private fun splitIntoNotCachedAndCached(sources: Set, + cache: Cache): Pair, List> { + if (sources.isEmpty()) return emptySet() to emptyList() + + val notCached = mutableSetOf() + val cached = mutableListOf() + + synchronized(cache) { + // despite cache is thread-safe, + // here we have sync block to guarantee only one reader + // and as a result not run processing for sources that are already evaluating + + sources.forEach { source -> + cache.getIfPresent(source).let { + if (it == null) { + notCached.add(source) + + // mark this source as evaluating + // if source processing failed, this value would mean that such source was handled + cache.put(source, StubPackagesForSource.EMPTY) + } + else { + cached.add(it) + } + } + } + } + + return notCached to cached + } + + private fun toRequirementsAndExtraArgs(cached: List, + ignoredStubPackages: List): Pair, List> { + if (cached.isEmpty()) return emptyList() to emptyList() + + val requirements = cached.asSequence() + .flatMap { it.packages.entries.asSequence() } + .filterNot { isIgnoredStubPackage(it.key, it.value.first, ignoredStubPackages) } + .map { + pyRequirement(it.key, PyRequirementRelation.EQ, it.value.first) + } + .toList() + if (requirements.isEmpty()) return emptyList() to emptyList() + + val args = sequenceOf("--no-deps") + + cached.asSequence().flatMap { pkgs -> pkgs.packages.values.asSequence().map { it.second }.flatten() } + return requirements to args.toList() + } + + private fun createInstallStubPackagesQuickFix(reqs: List, + args: List, + module: Module, + sdk: Sdk, + packageManager: PyPackageManager): LocalQuickFix { + val project = module.project + val stubPkgNamesToInstall = reqs.mapTo(mutableSetOf()) { it.name } + + val installationListener = object : PyPackageManagerUI.Listener { + override fun started() { + ServiceManager.getService(project, PyStubPackagesInstallingStatus::class.java).markAsInstalling(stubPkgNamesToInstall) } - status.unmarkAsInstalling(stubPkgNamesToInstall) + override fun finished(exceptions: MutableList?) { + val status = ServiceManager.getService(project, PyStubPackagesInstallingStatus::class.java) + + val stubPkgsToUninstall = PyStubPackagesCompatibilityInspection + .findIncompatibleRuntimeToStubPackages(sdk) { it.name in stubPkgNamesToInstall } + .map { it.second } + + if (stubPkgsToUninstall.isNotEmpty()) { + val stubPkgNamesToUninstall = stubPkgsToUninstall.mapTo(mutableSetOf()) { it.name } + + val uninstallationListener = object : PyPackageManagerUI.Listener { + override fun started() {} + + override fun finished(exceptions: MutableList?) { + status.unmarkAsInstalling(stubPkgNamesToUninstall) + val reqsToIgnore = stubPkgsToUninstall.map { pyRequirement(it.name, PyRequirementRelation.EQ, it.version) } + addStubPackagesToIgnore(reqsToIgnore, stubPkgNamesToUninstall, project, packageManager) + } + } + + val plural = stubPkgNamesToUninstall.size > 1 + val content = "Suggested ${stubPkgNamesToUninstall.joinToString { "'$it'" }} " + + "${if (plural) "are" else "is"} incompatible with your current environment.
" + + "${if (plural) "These" else "This"} stub package${if (plural) "s" else ""} will be removed and ignored until new version is released." + + BALLOON_NOTIFICATIONS.createNotification(content, NotificationType.WARNING).notify(project) + PyPackageManagerUI(project, sdk, uninstallationListener).uninstall(stubPkgsToUninstall) + + stubPkgNamesToInstall.removeAll(stubPkgNamesToUninstall) + } + + status.unmarkAsInstalling(stubPkgNamesToInstall) + } + } + + val name = "Install stub package" + if (reqs.size > 1) "s" else "" + return PyInstallRequirementsFix(name, module, sdk, reqs, args, installationListener) + } + + private fun createIgnorePackagesQuickFix(reqs: List, packageManager: PyPackageManager): LocalQuickFix { + return object : LocalQuickFix { + override fun getFamilyName() = "Ignore package" + if (reqs.size > 1) "s" else "" + + override fun applyFix(project: Project, descriptor: ProblemDescriptor) { + this@Visitor.addStubPackagesToIgnore(reqs, reqs.mapTo(mutableSetOf()) { it.name }, project, packageManager) + } } } - val name = "Install stub package" + if (reqs.size > 1) "s" else "" - return PyInstallRequirementsFix(name, module, sdk, reqs, args, installationListener) - } + private fun addStubPackagesToIgnore(stubPackages: List, + stubPackagesNames: Set, + project: Project, + packageManager: PyPackageManager) { + ignoredPackages.removeIf { packageManager.parseRequirement(it)?.name in stubPackagesNames } + ignoredPackages.addAll(stubPackages.map { it.presentableText }) - private fun createIgnorePackagesQuickFix(reqs: List, packageManager: PyPackageManager): LocalQuickFix { - return object : LocalQuickFix { - override fun getFamilyName() = "Ignore package" + if (reqs.size > 1) "s" else "" - - override fun applyFix(project: Project, descriptor: ProblemDescriptor) { - this@PyStubPackagesAdvertiser.addStubPackagesToIgnore(reqs, reqs.mapTo(mutableSetOf()) { it.name }, project, packageManager) - } + ProjectInspectionProfileManager.getInstance(project).fireProfileChanged() } - } - private fun addStubPackagesToIgnore(stubPackages: List, - stubPackagesNames: Set, - project: Project, - packageManager: PyPackageManager) { - ignoredPackages.removeIf { packageManager.parseRequirement(it)?.name in stubPackagesNames } - ignoredPackages.addAll(stubPackages.map { it.presentableText }) - - ProjectInspectionProfileManager.getInstance(project).fireProfileChanged() - } - - private fun isIgnoredStubPackage(name: String, version: String, ignoredStubPackages: List): Boolean { - val stubPackage = PyPackage(name, version, null, emptyList()) - return ignoredStubPackages.any { stubPackage.matches(it) } + private fun isIgnoredStubPackage(name: String, version: String, ignoredStubPackages: List): Boolean { + val stubPackage = PyPackage(name, version, null, emptyList()) + return ignoredStubPackages.any { stubPackage.matches(it) } + } } }