fix: fix interface error issues

unit-mesh · Jul 14, 2023 · 9f04d13 · 9f04d13
1 parent 0f17226
commit 9f04d13
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 4 deletions.
diff --git a/src/main/kotlin/cc/unitmesh/devti/flow/AutoDevFlow.kt b/src/main/kotlin/cc/unitmesh/devti/flow/AutoDevFlow.kt
@@ -178,6 +178,10 @@ class AutoDevFlow(
                 processor.createDto(code)
             }
 
+            processor.isRepository(code) -> {
+                processor.createRepository(code)
+            }
+
             else -> {
                 processor.createClass(code, null)
             }

diff --git a/src/main/kotlin/cc/unitmesh/devti/flow/JavaSpringBaseCrud.kt b/src/main/kotlin/cc/unitmesh/devti/flow/JavaSpringBaseCrud.kt
@@ -39,6 +39,7 @@ class JavaSpringBaseCrud(val project: Project) : SpringBaseCrud {
     override fun getAllEntityFiles(): List<PsiFile> = filterFilesByFunc(::entityFilter)
     override fun getAllDtoFiles(): List<PsiFile> = filterFilesByFunc(::dtoFilter)
     override fun getAllServiceFiles(): List<PsiFile> = filterFilesByFunc(::serviceFilter)
+    override fun getAllRepositoryFiles(): List<PsiFile> = filterFilesByFunc(::repositoryFilter)
 
     private fun filterFilesByFunc(filter: KFunction1<PsiClass, Boolean>): List<PsiFile> {
         return runReadAction {
@@ -141,6 +142,10 @@ class JavaSpringBaseCrud(val project: Project) : SpringBaseCrud {
         return createClassByCode(code, getAllDtoFiles())
     }
 
+    override fun createRepository(code: String): DtClass? {
+        return createClassByCode(code, getAllRepositoryFiles())
+    }
+
     override fun createService(code: String): DtClass? {
         return createClassByCode(code, getAllServiceFiles())
     }
@@ -171,11 +176,11 @@ class JavaSpringBaseCrud(val project: Project) : SpringBaseCrud {
             runWriteAction {
                 val newClass = psiElementFactory.createClassFromText(code, null)
 
-                val regex = Regex("public\\s+class\\s+(\\w+)")
+                val regex = Regex("public\\s+(class|interface)\\s+(\\w+)")
                 val matchResult = regex.find(code)
 
-                val className = if (matchResult?.groupValues?.get(1) != null) {
-                    matchResult.groupValues[1]
+                val className = if (matchResult?.groupValues?.get(2) != null) {
+                    matchResult.groupValues[2]
                 } else if (newClass.identifyingElement?.text != null) {
                     newClass.identifyingElement?.text
                 } else {

diff --git a/src/main/kotlin/cc/unitmesh/devti/flow/SpringBaseCrud.kt b/src/main/kotlin/cc/unitmesh/devti/flow/SpringBaseCrud.kt
@@ -17,13 +17,15 @@ interface SpringBaseCrud {
     fun getAllEntityFiles(): List<PsiFile>
     fun getAllDtoFiles(): List<PsiFile>
     fun getAllServiceFiles(): List<PsiFile>
+    fun getAllRepositoryFiles(): List<PsiFile>
 
 
     fun createControllerOrUpdateMethod(targetController: String, code: String, isControllerExist: Boolean)
     fun createController(endpoint: String, code: String): DtClass?
     fun createEntity(code: String): DtClass?
     fun createService(code: String): DtClass?
     fun createDto(code: String): DtClass?
+    fun createRepository(code: String): DtClass?
     fun createClass(code: String, packageName: String?): DtClass?
 
     fun dtoFilter(clazz: PsiClass): Boolean {
@@ -46,6 +48,11 @@ interface SpringBaseCrud {
             it == "org.springframework.stereotype.Service"
         }
 
+    fun repositoryFilter(clazz: PsiClass): Boolean = clazz.annotations
+        .map { it.qualifiedName }.any {
+            it == "org.springframework.stereotype.Repository"
+        }
+
     fun entityFilter(clazz: PsiClass): Boolean = clazz.annotations
         .map { it.qualifiedName }.any {
             it == "javax.persistence.Entity"
@@ -103,4 +110,17 @@ interface SpringBaseCrud {
         return regex.containsMatchIn(code)
     }
 
+    fun isRepository(code: String): Boolean {
+        if (code.contains("@Repository")) {
+            return true
+        }
+
+        if (code.contains("import org.springframework.stereotype.Repository")) {
+            return true
+        }
+
+        // regex to match `public class xxRepository`
+        val regex = Regex("public\\s+class\\s+\\w+Repository")
+        return regex.containsMatchIn(code)
+    }
 }
diff --git a/src/main/resources/prompts/openai/create_service_and_repository.txt b/src/main/resources/prompts/openai/create_service_and_repository.txt
@@ -22,7 +22,8 @@ public class {serviceName} {
 import org.springframework.data.repository.CrudRepository;
 import org.springframework.stereotype.Repository;
 
-public class {xxxRepository} extends JpaRepository<{xxx}, Long> {
+@Repository
+public interface {xxxRepository} extends JpaRepository<{xxx}, Long> {
 
 }
 ```