feat(kotlin): init code modifier

java/src/main/kotlin/cc/unitmesh/idea/context/JavaCodeModifier.kt

@@ -3,6 +3,7 @@ package cc.unitmesh.idea.context
 import cc.unitmesh.devti.context.builder.CodeModifier
 import cc.unitmesh.idea.service.JavaWriteTestService
 import com.intellij.lang.Language
+import com.intellij.lang.java.JavaLanguage
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.command.WriteCommandAction
@@ -15,19 +16,24 @@ import com.intellij.psi.PsiElementFactory
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiManager
 
-class JavaCodeModifier : CodeModifier {
+open class JavaCodeModifier : CodeModifier {
     companion object {
         val log = logger<JavaWriteTestService>()
     }
 
     override fun isApplicable(language: Language): Boolean {
-        return language == Language.findLanguageByID("JAVA")
+        return language is JavaLanguage
     }
 
+    fun lookupFile(
+        project: Project,
+        sourceFile: VirtualFile
+    ) = PsiManager.getInstance(project).findFile(sourceFile) as PsiJavaFile
+
     override fun insertTestCode(sourceFile: VirtualFile, project: Project, code: String): Boolean {
-        JavaWriteTestService.log.info("methodCode: $code")
+        log.info("methodCode: $code")
         if (!code.contains("@Test")) {
-            JavaWriteTestService.log.error("methodCode does not contain @Test annotation: $code")
+            log.error("methodCode does not contain @Test annotation: $code")
             return false
         }
 
@@ -42,10 +48,10 @@ class JavaCodeModifier : CodeModifier {
     override fun insertMethod(sourceFile: VirtualFile, project: Project, code: String): Boolean {
         ApplicationManager.getApplication().invokeLater {
             val rootElement = runReadAction {
-                val psiJavaFile = PsiManager.getInstance(project).findFile(sourceFile) as PsiJavaFile
+                val psiJavaFile = lookupFile(project, sourceFile)
                 val psiClass = psiJavaFile.classes.firstOrNull()
                 if (psiClass == null) {
-                    JavaWriteTestService.log.error("Failed to find PsiClass in the source file: $psiJavaFile, code: $code")
+                    log.error("Failed to find PsiClass in the source file: $psiJavaFile, code: $code")
                     return@runReadAction null
                 }
 
@@ -56,10 +62,10 @@ class JavaCodeModifier : CodeModifier {
 
             val newTestMethod = psiElementFactory.createMethodFromText(code, rootElement)
             if (rootElement.findMethodsByName(newTestMethod.name, false).isNotEmpty()) {
-                JavaWriteTestService.log.error("Method already exists in the class: ${newTestMethod.name}")
+                log.error("Method already exists in the class: ${newTestMethod.name}")
             }
 
-            JavaWriteTestService.log.info("newTestMethod: ${newTestMethod.text}")
+            log.info("newTestMethod: ${newTestMethod.text}")
 
             WriteCommandAction.runWriteCommandAction(project) {
                 val lastMethod = rootElement.methods.lastOrNull()
@@ -80,7 +86,7 @@ class JavaCodeModifier : CodeModifier {
     override fun insertClass(sourceFile: VirtualFile, project: Project, code: String): Boolean {
         log.info("start insertClassCode: $code")
         WriteCommandAction.runWriteCommandAction(project) {
-            val psiFile = PsiManager.getInstance(project).findFile(sourceFile) as PsiJavaFile
+            val psiFile = lookupFile(project, sourceFile)
             val document = psiFile.viewProvider.document!!
             document.insertString(document.textLength, code)
         }

kotlin/src/main/kotlin/cc/unitmesh/kotlin/context/KotlinCodeModifier.kt

@@ -0,0 +1,93 @@
+package cc.unitmesh.kotlin.context
+
+import cc.unitmesh.devti.context.builder.CodeModifier
+import com.intellij.lang.Language
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.project.guessProjectDir
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiManager
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtPsiFactory
+
+class KotlinCodeModifier : CodeModifier {
+    companion object {
+        val log = logger<KotlinCodeModifier>()
+    }
+
+    override fun isApplicable(language: Language): Boolean {
+        return language is KotlinLanguage
+    }
+
+    fun lookupFile(
+        project: Project,
+        sourceFile: VirtualFile
+    ) = PsiManager.getInstance(project).findFile(sourceFile) as KtFile
+
+    override fun insertTestCode(sourceFile: VirtualFile, project: Project, code: String): Boolean {
+        log.info("methodCode: $code")
+        if (!code.contains("@Test")) {
+            log.error("methodCode does not contain @Test annotation: $code")
+            return false
+        }
+
+        if (code.startsWith("import") && code.contains("class ")) {
+            return insertClass(sourceFile, project, code)
+        }
+
+        insertMethod(sourceFile, project, code)
+        return true
+    }
+
+    override fun insertMethod(sourceFile: VirtualFile, project: Project, code: String): Boolean {
+        ApplicationManager.getApplication().invokeLater {
+            val rootElement = runReadAction {
+                val ktFile = lookupFile(project, sourceFile)
+                val psiClass = ktFile.classes.firstOrNull()
+                if (psiClass == null) {
+                    log.error("Failed to find PsiClass in the source file: $ktFile, code: $code")
+                    return@runReadAction null
+                }
+
+                return@runReadAction psiClass
+            } ?: return@invokeLater
+
+
+            val newTestMethod =  KtPsiFactory(project).createFunction(code)
+            if (rootElement.findMethodsByName(newTestMethod.name, false).isNotEmpty()) {
+                log.error("Method already exists in the class: ${newTestMethod.name}")
+            }
+
+            log.info("newTestMethod: ${newTestMethod.text}")
+
+            WriteCommandAction.runWriteCommandAction(project) {
+                val lastMethod = rootElement.methods.lastOrNull()
+                val lastMethodEndOffset = lastMethod?.textRange?.endOffset ?: 0
+
+                val document = PsiDocumentManager.getInstance(project).getDocument(rootElement.containingFile)
+                document?.insertString(lastMethodEndOffset, "\n    ")
+                document?.insertString(lastMethodEndOffset, newTestMethod.text)
+            }
+
+            project.guessProjectDir()?.refresh(true, true)
+        }
+
+        return true
+    }
+
+    override fun insertClass(sourceFile: VirtualFile, project: Project, code: String): Boolean {
+        log.info("start insertClassCode: $code")
+        WriteCommandAction.runWriteCommandAction(project) {
+            val psiFile = lookupFile(project, sourceFile)
+            val document = psiFile.viewProvider.document!!
+            document.insertString(document.textLength, code)
+        }
+
+        return true
+    }
+}

kotlin/src/main/resources/cc.unitmesh.kotlin.xml

@@ -6,6 +6,9 @@
     </dependencies>
 
     <extensions defaultExtensionNs="cc.unitmesh">
+        <codeModifier language="kotlin"
+                      implementationClass="cc.unitmesh.kotlin.context.KotlinCodeModifier"/>
+
         <chatContextProvider implementation="cc.unitmesh.kotlin.provider.KotlinTestContextProvider"/>
 
         <classContextBuilder language="kotlin"