feat: check for issue

unit-mesh · Dec 8, 2023 · 5ee0c26 · 5ee0c26
1 parent 32ddc27
commit 5ee0c26
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 8 deletions.
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/InstructionBuilder.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/InstructionBuilder.kt
@@ -1,9 +1,13 @@
 package cc.unitmesh.pick.prompt
 
 import cc.unitmesh.pick.picker.InstructionJob
+import cc.unitmesh.quality.CodeQualityType
+import cc.unitmesh.quality.QualityAnalyser
+import chapi.domain.core.CodeDataStruct
 
 data class InstructionContext(
     val job: InstructionJob,
+    val qualityTypes: List<CodeQualityType>,
     val fileTree: HashMap<String, InstructionJob>,
 )
 
@@ -54,15 +58,22 @@ interface InstructionBuilder<T> {
      */
     fun build(): List<Instruction>
 
+    fun hasIssue(node: CodeDataStruct, types: List<CodeQualityType>): Boolean {
+        return QualityAnalyser.create(types).map { analyser ->
+            analyser.analysis(listOf(node)).isNotEmpty()
+        }.any { it }
+    }
+
     companion object {
         fun build(
-            types: List<InstructionType>,
+            instructionTypes: List<InstructionType>,
+            qualityTypes: List<CodeQualityType>,
             fileTree: HashMap<String, InstructionJob>,
             job: InstructionJob,
         ): Collection<Instruction> {
-            val instructionContext = InstructionContext(job, fileTree)
+            val instructionContext = InstructionContext(job, qualityTypes, fileTree)
 
-            return types.map { type ->
+            return instructionTypes.map { type ->
                 type.builder(instructionContext).build()
             }.flatten()
         }

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/builder/RelatedCodeCompletionBuilder.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/prompt/builder/RelatedCodeCompletionBuilder.kt
@@ -23,9 +23,14 @@ class RelatedCodeCompletionBuilder(private val context: InstructionContext) :
         val container = context.job.container ?: return emptyList()
 
         // 1. collection all related data structure by imports if exists in a file tree
-        val relatedDataStructure = container.Imports.mapNotNull {
-            context.fileTree[it.Source]?.container?.DataStructures
-        }.flatten()
+        val relatedDataStructure = container.Imports
+            .mapNotNull {
+                context.fileTree[it.Source]?.container?.DataStructures
+            }
+            .flatten()
+            .filter {
+                hasIssue(it, context.qualityTypes)
+            }
 
         // 2. convert all related data structure to uml
         val relatedCode = relatedDataStructure.joinToString("\n", transform = CodeDataStruct::toUml)

diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerContext.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/WorkerContext.kt
@@ -1,7 +1,9 @@
 package cc.unitmesh.pick.worker
 
 import cc.unitmesh.pick.prompt.InstructionType
+import cc.unitmesh.quality.CodeQualityType
 
 data class WorkerContext(
-    val builderType: List<InstructionType>,
+    val builderTypes: List<InstructionType>,
+    val qualityTypes: List<CodeQualityType>
 )
diff --git a/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/worker/JavaLangWorker.kt b/unit-picker/src/main/kotlin/cc/unitmesh/pick/worker/worker/JavaLangWorker.kt
@@ -61,7 +61,7 @@ class JavaLangWorker(val workerContext: WorkerContext) : LangWorker() {
 
     override suspend fun start(): Collection<Instruction> = coroutineScope {
         return@coroutineScope jobs.map {
-            InstructionBuilder.build(workerContext.builderType, fileTree, it)
+            InstructionBuilder.build(workerContext.builderTypes, workerContext.qualityTypes, fileTree, it)
         }.flatten()
     }