feat: add import optimizer

Author: koxudaxi

resources/META-INF/plugin.xml

@@ -17,6 +17,7 @@
                          implementationClass="com.koxudaxi.ruff.RuffInspection"/>
         <externalAnnotator language="Python" implementationClass="com.koxudaxi.ruff.RuffExternalAnnotator"/>
         <formattingService implementation="com.koxudaxi.ruff.RuffAsyncFormatter" order="first"/>
+        <lang.importOptimizer language="Python" implementationClass="com.koxudaxi.ruff.RuffImportOptimizer" order="first"/>
         <platform.backend.documentation.targetProvider
                 implementation="com.koxudaxi.ruff.RuffNoqaDocumentationTargetProvider"/>
         <actionOnSave id="RuffFormatterActionOnSave" implementation="com.koxudaxi.ruff.RuffActionOnSave" order="last"/>

src/com/koxudaxi/ruff/Ruff.kt

@@ -92,8 +92,17 @@ val NO_FIX_OUTPUT_FORMAT_ARGS = ARGS_BASE + listOf("--no-fix", "--output-format"
 val FORMAT_ARGS = listOf("format", "--force-exclude", "--quiet")
 val FORMAT_CHECK_ARGS = FORMAT_ARGS + listOf("--check")
 val FORMAT_RANGE_ARGS = listOf("--range")
+val SELECT_ARGS = listOf("--select")
+val IMPORT_SORT_RULES = listOf("I")
+val REMOVE_UNUSED_IMPORTS_RULES = listOf("F401")
+val OPTIMIZE_IMPORTS_RULES = SELECT_ARGS + listOf((IMPORT_SORT_RULES + REMOVE_UNUSED_IMPORTS_RULES).joinToString(","))
 val LSP_ARGS_BASE = listOf("server")
 val PREVIEW_ARGS = listOf("--preview")
+val Project.OPTIMIZE_IMPORTS_ARGS: List<String>
+    get() = when {
+        RuffCacheService.hasCheck(this) == true -> CHECK + FIX_ARGS_BASE + OPTIMIZE_IMPORTS_RULES
+        else -> FIX_ARGS_BASE + OPTIMIZE_IMPORTS_RULES
+    }
 val Project.LSP_ARGS: List<String>
     get() = when {
         RuffCacheService.hasStableServer(this) == true -> LSP_ARGS_BASE
@@ -445,20 +454,22 @@ data class SourceFile(
 ) {
     val text: String? by lazy {
         val text = when (reloadText) {
-            true -> documentationManager.getDocument(psiFile)?.text
+            true -> currentText
             else -> psiFile.text
         } ?: return@lazy null
         if (textRange == null) return@lazy text
         if (textRange.endOffset <= text.length) textRange.substring(text)
         text.substring(textRange.startOffset, text.length)
     }
+    val currentText: String? get() = documentationManager.getDocument(psiFile)?.text
     private val documentationManager: PsiDocumentManager get() = PsiDocumentManager.getInstance(project)
     val project: Project get() = psiFile.project
     val virtualFile: VirtualFile? get() = psiFile.virtualFile
 
     val name: String get() = psiFile.name
-    val asStdin: ByteArray? get() = text?.toCharArray()?.toByteArrayAndClear()
-
+    val asStdin: ByteArray? get() = text?.let {convertToBytes(it) }
+    val asCurrentTextStdin: ByteArray? get() = currentText?.let {convertToBytes(it) }
+    private fun convertToBytes(text: String): ByteArray = text.toCharArray().toByteArrayAndClear()
     fun hasSameContentAsDocument(document: Document): Boolean = document.charsSequence.contentEquals(text)
 }
 

src/com/koxudaxi/ruff/RuffAsyncFormatter.kt

@@ -4,14 +4,16 @@ import com.intellij.formatting.FormattingContext
 import com.intellij.formatting.service.AsyncDocumentFormattingService
 import com.intellij.formatting.service.AsyncFormattingRequest
 import com.intellij.formatting.service.FormattingService
+import com.intellij.lang.ImportOptimizer
 import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.psi.PsiFile
 import com.jetbrains.python.packaging.PyExecutionException
 import java.io.FileNotFoundException
 
 
 class RuffAsyncFormatter : AsyncDocumentFormattingService() {
-    private val FEATURES: Set<FormattingService.Feature> = setOf(FormattingService.Feature.FORMAT_FRAGMENTS)
+    private val FEATURES: Set<FormattingService.Feature> =
+        setOf(FormattingService.Feature.FORMAT_FRAGMENTS, FormattingService.Feature.OPTIMIZE_IMPORTS)
 
     override fun getFeatures(): Set<FormattingService.Feature> {
         return FEATURES
@@ -60,7 +62,8 @@ class RuffAsyncFormatter : AsyncDocumentFormattingService() {
                             false
                         )
                             ?: return@runCatching null
-                        val formatCommandStdout = runRuff(formatCommandArgs, currentText.toByteArray()) ?: return@runCatching null
+                        val formatCommandStdout =
+                            runRuff(formatCommandArgs, currentText.toByteArray()) ?: return@runCatching null
                         return@runCatching assertResult(currentText, formatCommandStdout)
                     }
 
@@ -110,4 +113,14 @@ class RuffAsyncFormatter : AsyncDocumentFormattingService() {
     override fun getName(): String {
         return "Ruff Formatter"
     }
+
+    override fun getImportOptimizers(file: PsiFile): Set<ImportOptimizer?> {
+        val importOptimizer = RuffImportOptimizer()
+        return if (importOptimizer.supports(file)) {
+            setOf(importOptimizer)
+        } else {
+            emptySet()
+        }
+    }
+
 }

src/com/koxudaxi/ruff/RuffImportOptimizer.kt

@@ -0,0 +1,47 @@
+package com.koxudaxi.ruff
+
+import com.intellij.lang.ImportOptimizer
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.command.CommandProcessor
+import com.intellij.psi.PsiFile
+import com.jetbrains.python.psi.PyUtil
+
+class RuffImportOptimizer : ImportOptimizer {
+    override fun supports(psiFile: PsiFile): Boolean = psiFile.isApplicableTo
+
+    override fun processFile(psiFile: PsiFile): Runnable {
+        val project = psiFile.project
+        val sourceFile = psiFile.sourceFile
+        val optimizeImportsArgs = project.OPTIMIZE_IMPORTS_ARGS
+
+        return Runnable {
+            val stdin = sourceFile.asCurrentTextStdin ?: return@Runnable
+            runRuffInBackground(project, stdin, project.OPTIMIZE_IMPORTS_ARGS, "Optimize Imports with Ruff") {
+                val result = executeOnPooledThread(project, null as String?) {
+                    runRuff(sourceFile, optimizeImportsArgs)
+                }
+
+                if (result != null && result.isNotBlank()) {
+                    ApplicationManager.getApplication().invokeLater {
+                        CommandProcessor.getInstance().executeCommand(
+                            project,
+                            {
+                                ApplicationManager.getApplication().runWriteAction {
+                                    PyUtil.updateDocumentUnblockedAndCommitted(
+                                        psiFile
+                                    ) { document ->
+                                        if (document.text != result) {
+                                            document.setText(result)
+                                        }
+                                    }
+                                }
+                            },
+                            "Ruff: Optimize Imports",
+                            null
+                        )
+                    }
+                }
+            }
+        }
+    }
+}