feat: add condition for check method

unit-mesh · Jul 13, 2023 · ea6c87a · ea6c87a
1 parent 631c818
commit ea6c87a
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 11 deletions.
diff --git a/src/main/kotlin/cc/unitmesh/devti/actions/chat/CodeCompleteAction.kt b/src/main/kotlin/cc/unitmesh/devti/actions/chat/CodeCompleteAction.kt
@@ -16,8 +16,13 @@ class CodeCompleteAction : ChatBaseAction() {
         val document = editor.document
 
         val primaryCaret = editor.caretModel.primaryCaret;
-        val start = primaryCaret.selectionStart;
-        val end = primaryCaret.selectionEnd
+        val start = primaryCaret.selectionStart
+        var end = primaryCaret.selectionEnd
+
+        if (start == end) {
+            primaryCaret.setSelection(start, start + 1)
+            end = start + 1
+        }
 
         return { response ->
             WriteCommandAction.runWriteCommandAction(project) {

diff --git a/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingService.kt b/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingService.kt
@@ -2,7 +2,7 @@ package cc.unitmesh.devti.gui.chat
 
 import cc.unitmesh.devti.AutoDevBundle
 import cc.unitmesh.devti.connector.ConnectorService
-import cc.unitmesh.devti.parser.CodePostProcessor
+import cc.unitmesh.devti.parser.JavaCodePostProcessor
 import com.intellij.openapi.application.ApplicationManager
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.runBlocking
@@ -64,7 +64,7 @@ class ChatCodingService(var actionType: ChatBotActionType) {
 
         if (match != null) return match.groupValues[1].trim()
 
-        return CodePostProcessor(prefixText, suffixText, content).execute()
+        return JavaCodePostProcessor(prefixText, suffixText, content).execute()
     }
 }
 

diff --git a/...nitmesh/devti/parser/CodePostProcessor.kt → ...esh/devti/parser/JavaCodePostProcessor.kt b/...nitmesh/devti/parser/CodePostProcessor.kt → ...esh/devti/parser/JavaCodePostProcessor.kt
@@ -3,10 +3,10 @@ package cc.unitmesh.devti.parser
 /**
  * will format complete code by prefix code and suffix code
  */
-class CodePostProcessor(
-    val prefixCode: String,
-    val suffixCode: String,
-    val completeCode: String
+class JavaCodePostProcessor(
+    private val prefixCode: String,
+    private val suffixCode: String,
+    private val completeCode: String
 ) {
     private val spaceRegex = Regex("\\s+")
 
@@ -36,6 +36,11 @@ class CodePostProcessor(
             result = result.split("\n").joinToString("\n") { "    $it" }
         }
 
+        // if suffix ends with "}", and complete code ends with "}\n}", then remove complete code's last "}"
+        if (result.endsWith("}\n}") and suffixCode.endsWith("}")) {
+            result = result.substring(0, result.length - 1)
+        }
+
         // if lastLineSpaceCount > 0, then remove same space in result begin if exists
         if (lastLineSpaceCount > 0) {
             val spaceRegex = Regex("^\\s{$lastLineSpaceCount}")

diff --git a/...esh/devti/parser/CodePostProcessorTest.kt → ...devti/parser/JavaCodePostProcessorTest.kt b/...esh/devti/parser/CodePostProcessorTest.kt → ...devti/parser/JavaCodePostProcessorTest.kt
@@ -3,7 +3,7 @@ package cc.unitmesh.devti.parser
 import org.junit.Assert.assertEquals
 import org.junit.Test
 
-class CodePostProcessorTest {
+class JavaCodePostProcessorTest {
     @Test
     fun should_handle_with_no_indent_code() {
         val prefix = "public class Test {\n"
@@ -15,7 +15,7 @@ public void test() {
 }
         """.trimIndent()
 
-        val postProcessor = CodePostProcessor(prefix, suffix, complete)
+        val postProcessor = JavaCodePostProcessor(prefix, suffix, complete)
         val result = postProcessor.execute()
 
         assertEquals(
@@ -33,7 +33,7 @@ public void test() {
         val complete = """public void test() {
 }
 }"""
-        val postProcessor = CodePostProcessor(prefix, suffix, complete)
+        val postProcessor = JavaCodePostProcessor(prefix, suffix, complete)
         val result = postProcessor.execute()
 
         assertEquals(