fix: fix kotlin code modifier issue

unit-mesh · Dec 18, 2023 · 501fccc · 501fccc
1 parent 96b13cb
commit 501fccc
Showing 1 changed file with 21 additions and 12 deletions.
diff --git a/kotlin/src/main/kotlin/cc/unitmesh/kotlin/provider/KotlinCodeModifier.kt b/kotlin/src/main/kotlin/cc/unitmesh/kotlin/provider/KotlinCodeModifier.kt
@@ -56,21 +56,30 @@ class KotlinCodeModifier : CodeModifier {
                 return@runReadAction psiClass
             } ?: return@invokeLater
 
+            try {
+                val newTestMethod = KtPsiFactory(project).createFunction(code)
+                if (rootElement.findMethodsByName(newTestMethod.name, false).isNotEmpty()) {
+                    log.error("Method already exists in the class: ${newTestMethod.name}")
+                }
 
-            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}")
+                log.info("newTestMethod: ${newTestMethod.text}")
 
-            WriteCommandAction.runWriteCommandAction(project) {
-                val lastMethod = rootElement.methods.lastOrNull()
-                val lastMethodEndOffset = lastMethod?.textRange?.endOffset ?: 0
+                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)
+                    val document = PsiDocumentManager.getInstance(project).getDocument(rootElement.containingFile)
+                    document?.insertString(lastMethodEndOffset, "\n    ")
+                    document?.insertString(lastMethodEndOffset, newTestMethod.text)
+                }
+            } catch (e: Exception) {
+                log.warn("Failed to insert method: $code", e)
+                // append to the end of the file
+                WriteCommandAction.runWriteCommandAction(project) {
+                    val document = PsiDocumentManager.getInstance(project).getDocument(rootElement.containingFile)!!
+                    document?.insertString(document.textLength, "\n    ")
+                    document?.insertString(document.textLength, code)
+                }
             }
 
             project.guessProjectDir()?.refresh(true, true)