diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/HuggingFaceRelevantLibraries.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/HuggingFaceRelevantLibraries.kt deleted file mode 100644 index 7ca41cefe1aa..000000000000 --- a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/HuggingFaceRelevantLibraries.kt +++ /dev/null @@ -1,10 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.community.impl.huggingFace - -object HuggingFaceRelevantLibraries { - val relevantLibraries = setOf( - "diffusers", "transformers", "allennlp", "spacy", - "asteroid", "flair", "keras", "sentence-transformers", - "stable-baselines3", "adapters", "huggingface_hub" - ) -} diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceImportedLibrariesManager.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceImportedLibrariesManager.kt deleted file mode 100644 index 786a45f85658..000000000000 --- a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceImportedLibrariesManager.kt +++ /dev/null @@ -1,154 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.community.impl.huggingFace.service - -import com.intellij.openapi.Disposable -import com.intellij.openapi.components.Service -import com.intellij.openapi.diagnostic.thisLogger -import com.intellij.openapi.editor.EditorFactory -import com.intellij.openapi.editor.event.DocumentEvent -import com.intellij.openapi.editor.event.DocumentListener -import com.intellij.openapi.fileEditor.FileDocumentManager -import com.intellij.openapi.project.Project -import com.intellij.openapi.roots.ProjectFileIndex -import com.intellij.openapi.roots.ProjectRootManager -import com.intellij.openapi.vfs.VirtualFileManager -import com.intellij.openapi.vfs.newvfs.BulkFileListener -import com.intellij.openapi.vfs.newvfs.events.VFileEvent -import com.intellij.psi.PsiManager -import com.intellij.python.community.impl.huggingFace.HuggingFaceRelevantLibraries -import com.intellij.python.community.impl.huggingFace.cache.HuggingFaceCacheFillService -import com.intellij.util.concurrency.annotations.RequiresBackgroundThread -import com.intellij.util.messages.MessageBusConnection -import com.jetbrains.python.psi.PyFile -import com.jetbrains.python.psi.PyImportElement -import org.jetbrains.annotations.ApiStatus - -@ApiStatus.Internal -@Service(Service.Level.PROJECT) -class HuggingFaceImportedLibrariesManager(val project: Project) : Disposable { - private var libraryImportStatus: LibraryImportStatus = LibraryImportStatus.NOT_CHECKED - private var cacheTimestamp: Long = 0 - private val connection: MessageBusConnection = project.messageBus.connect(this) - private var documentListener: DocumentListener? = null - private val cacheFillService: HuggingFaceCacheFillService = project.getService(HuggingFaceCacheFillService::class.java) - private val librariesChecker = HuggingFaceLibraryImportChecker(project) - private enum class LibraryImportStatus { NOT_CHECKED, IMPORTED, NOT_IMPORTED } - - init { setupListeners() } - - private fun setupListeners() { - connection.subscribe(VirtualFileManager.VFS_CHANGES, HuggingFaceFileChangesListener { checkLibraryImportStatusInProject() }) - val documentListener = HuggingFaceImportDetectionListener(project) { pyFile -> checkLibraryImportStatusInFile(pyFile) } - EditorFactory.getInstance().eventMulticaster.addDocumentListener(documentListener, connection) - this.documentListener = documentListener - } - - @RequiresBackgroundThread - fun isLibraryImported(): Boolean { - if (libraryImportStatus != LibraryImportStatus.IMPORTED) checkLibraryImportStatusInProject() - return libraryImportStatus == LibraryImportStatus.IMPORTED - } - - private fun checkLibraryImportStatusInFile(pyFile: PyFile) { - val isImported = librariesChecker.isAnyHFLibraryImportedInFile(pyFile) - updateLibraryImportStatus(isImported) - } - - private fun checkLibraryImportStatusInProject() { - if (libraryImportStatus == LibraryImportStatus.IMPORTED) return - val isUpdateTime = System.currentTimeMillis() - cacheTimestamp > HuggingFaceLibrariesManagerConfig.INVALIDATION_THRESHOLD_MS - if (libraryImportStatus == LibraryImportStatus.NOT_CHECKED || isUpdateTime) - { - val isImported = librariesChecker.isAnyHFLibraryImportedInProject() - updateLibraryImportStatus(isImported) - cacheTimestamp = System.currentTimeMillis() - } - } - - private fun updateLibraryImportStatus(newStatus: Boolean) { - libraryImportStatus = if (newStatus) { - cacheFillService.triggerCacheFillIfNeeded() - detachListeners() - LibraryImportStatus.IMPORTED - } else { - LibraryImportStatus.NOT_IMPORTED - } - } - - private fun detachListeners() { - documentListener?.let { listener -> - EditorFactory.getInstance().eventMulticaster.removeDocumentListener(listener) - documentListener = null - } - connection.disconnect() - } - - override fun dispose() { - detachListeners() - } -} - -private class HuggingFaceLibraryImportChecker(val project: Project) { - fun isAnyHFLibraryImportedInProject(): Boolean { - var isLibraryImported = false - - ProjectFileIndex.getInstance(project).iterateContent { virtualFile -> - if (virtualFile.extension in listOf("py", "ipynb")) { - val pythonFile = PsiManager.getInstance(project).findFile(virtualFile) - if (pythonFile is PyFile) isLibraryImported = isLibraryImported or isAnyHFLibraryImportedInFile(pythonFile) - } - !isLibraryImported - } - - return isLibraryImported - } - - fun isAnyHFLibraryImportedInFile(file: PyFile): Boolean { - val isDirectlyImported = file.importTargets.any { importStmt -> - HuggingFaceRelevantLibraries.relevantLibraries.any { lib -> importStmt.importedQName.toString().contains(lib) } - } - - val isFromImported = file.fromImports.any { fromImport -> - HuggingFaceRelevantLibraries.relevantLibraries.any { lib -> fromImport.importSourceQName?.toString()?.contains(lib) == true } - } - - val isQualifiedImported: Boolean = file.importTargets.any { importStmt: PyImportElement? -> - HuggingFaceRelevantLibraries.relevantLibraries.any { lib: String -> importStmt?.importedQName?.components?.contains(lib) == true } - } - return isDirectlyImported || isFromImported || isQualifiedImported - } -} - -private class HuggingFaceFileChangesListener(private val onThresholdReached: () -> Unit) : BulkFileListener { - private var fileChangesCounter = 0 - - override fun after(events: List) { - try { - if (events.any { it.file?.extension in listOf("py", "ipynb") }) { - fileChangesCounter++ - if (fileChangesCounter >= HuggingFaceLibrariesManagerConfig.CHANGES_NUM_THRESHOLD) onThresholdReached() - } - } catch (e: Exception) { - thisLogger().warn("Exception in HuggingFaceFileChangesListener.after", e) - } - } -} - -private class HuggingFaceImportDetectionListener( - private val project: Project, - private val onImportDetected: (PyFile) -> Unit -) : DocumentListener { - override fun documentChanged(event: DocumentEvent) { - if (!event.newFragment.toString().contains("import")) return - val file = FileDocumentManager.getInstance().getFile(event.document) ?: return - if (!ProjectRootManager.getInstance(project).fileIndex.isInContent(file)) return - PsiManager.getInstance(project).findFile(file)?.let { psiFile -> - if (psiFile is PyFile) onImportDetected(psiFile) - } - } -} - -private object HuggingFaceLibrariesManagerConfig { - const val CHANGES_NUM_THRESHOLD = 10 - const val INVALIDATION_THRESHOLD_MS = 5 * 60 * 1000 -} diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceLibrariesTracker.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceLibrariesTracker.kt new file mode 100644 index 000000000000..8cf6fe563573 --- /dev/null +++ b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFaceLibrariesTracker.kt @@ -0,0 +1,80 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.python.community.impl.huggingFace.service + +import com.intellij.openapi.Disposable +import com.intellij.openapi.components.Service +import com.intellij.openapi.project.Project +import com.intellij.openapi.project.modules +import com.intellij.openapi.projectRoots.Sdk +import com.intellij.python.community.impl.huggingFace.cache.HuggingFaceCacheFillService +import com.intellij.util.messages.MessageBusConnection +import com.jetbrains.python.packaging.PyPackageInstallUtils +import com.jetbrains.python.packaging.common.PythonPackageManagementListener +import com.jetbrains.python.packaging.management.PythonPackageManager +import com.jetbrains.python.sdk.PythonSdkUtil +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.launch +import org.jetbrains.annotations.ApiStatus + +@ApiStatus.Internal +@Service(Service.Level.PROJECT) +class HuggingFaceLibrariesTracker( + private val project: Project, + private val coroutineScope: CoroutineScope +) : Disposable { + @Volatile private var isAnyHFLibraryInstalled: Boolean = false + private var connection: MessageBusConnection? = project.messageBus.connect(this) + private val cacheFillService: HuggingFaceCacheFillService = project.getService(HuggingFaceCacheFillService::class.java) + + private val relevantLibraries = setOf( + "diffusers", "transformers", "allennlp", "spacy", + "asteroid", "flair", "keras", "sentence-transformers", + "stable-baselines3", "adapters", "huggingface_hub" + ) + + init { + setupSdkListener() + } + + fun isAnyHFLibraryInstalled(): Boolean = isAnyHFLibraryInstalled + + private fun setupSdkListener() { + connection?.subscribe(PythonPackageManager.PACKAGE_MANAGEMENT_TOPIC, object : PythonPackageManagementListener { + override fun packagesChanged(sdk: Sdk) { + val projectSdk = getProjectPythonSdk() + + if (sdk == projectSdk) { + coroutineScope.launch(Dispatchers.IO) { + updateHFLibraryInstallStatus() + } + } + } + }) + } + + private fun detachSdkListener() { + connection?.disconnect() + connection = null + } + + private fun getProjectPythonSdk(): Sdk? = PythonSdkUtil.findPythonSdk(project.modules.firstOrNull()) + + private fun updateHFLibraryInstallStatus() { + if (isAnyHFLibraryInstalled) return // assuming that if was found once - always relevant + + val sdk = getProjectPythonSdk() ?: return + + if (isAnyHFLibraryInstalledInSdk(sdk)) { + isAnyHFLibraryInstalled = true + cacheFillService.triggerCacheFillIfNeeded() + detachSdkListener() + } + } + + private fun isAnyHFLibraryInstalledInSdk(sdk: Sdk): Boolean = relevantLibraries.any { lib -> + PyPackageInstallUtils.getPackageVersion(project, sdk, lib) != null + } + + override fun dispose() = detachSdkListener() +} \ No newline at end of file diff --git a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFacePluginManager.kt b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFacePluginManager.kt index 42724765b9b6..bc49477c0f11 100644 --- a/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFacePluginManager.kt +++ b/python/huggingFace/src/com/intellij/python/community/impl/huggingFace/service/HuggingFacePluginManager.kt @@ -12,8 +12,8 @@ import org.jetbrains.annotations.ApiStatus @ApiStatus.Internal @Service(Service.Level.PROJECT) class HuggingFacePluginManager(val project: Project) : Disposable { - private var libraryStatusChecker: HuggingFaceImportedLibrariesManager = project.getService(HuggingFaceImportedLibrariesManager::class.java) + private var libraryStatusChecker: HuggingFaceLibrariesTracker = project.getService(HuggingFaceLibrariesTracker::class.java) init { project.getService(HuggingFaceCacheUpdateHandler::class.java) } - fun isActive(): Boolean = libraryStatusChecker.isLibraryImported() && Registry.`is`("python.enable.hugging.face.cards") + fun isActive(): Boolean = libraryStatusChecker.isAnyHFLibraryInstalled() && Registry.`is`("python.enable.hugging.face.cards") override fun dispose() = Disposer.dispose(libraryStatusChecker) }