feat: init render option for children #52

unit-mesh · Nov 20, 2023 · 610e33b · 610e33b
1 parent 8ddc6ab
commit 610e33b
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 5 deletions.
diff --git a/java/src/main/kotlin/cc/unitmesh/idea/prompting/JavaContextPrompter.kt b/java/src/main/kotlin/cc/unitmesh/idea/prompting/JavaContextPrompter.kt
@@ -135,9 +135,10 @@ open class JavaContextPrompter : ContextPrompter() {
         val methodContext = MethodContextProvider(false, false).from(element)
         selectedText = methodContext.text
 
-        val datastructures = methodContext.inputOutputString()
+        val datastructures = methodContext.inputOutputString(true)
         additionContext += """
-input and output's Class/DataStructures: 
+// input and output's Class/DataStructures:
+
 $datastructures
 """
     }

diff --git a/src/main/kotlin/cc/unitmesh/devti/context/MethodContext.kt b/src/main/kotlin/cc/unitmesh/devti/context/MethodContext.kt
@@ -5,6 +5,7 @@ import cc.unitmesh.devti.isInProject
 import com.intellij.lang.LanguageCommenters
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiField
 import com.intellij.psi.PsiReference
 
 class MethodContext(
@@ -56,10 +57,17 @@ fun signature: ${signature ?: "_"}
         return query
     }
 
-    fun inputOutputString(): String {
+    fun inputOutputString(withChild: Boolean = false): String {
+        if (fanInOut.isEmpty()) return ""
+
+        // !!!
+        if (withChild) {
+            return inputOutputStringWithChild()
+        }
+
         var result = ""
         this.fanInOut.forEach {
-            val context = ClassContextProvider(false).from(it)
+            val context: ClassContext = ClassContextProvider(false).from(it)
             val element = context.root
 
             if (!isInProject(element.containingFile?.virtualFile!!, project)) {
@@ -77,4 +85,45 @@ fun signature: ${signature ?: "_"}
 
         return "```uml\n$result\n```\n"
     }
+
+    val nestedNodes: MutableMap<String, ClassContext> = mutableMapOf()
+    private fun inputOutputStringWithChild(): String {
+        val result = StringBuilder()
+        this.fanInOut.forEach {
+            val context: ClassContext = ClassContextProvider(false).from(it)
+            val element = context.root
+
+            if (!isInProject(element.containingFile?.virtualFile!!, project)) {
+                return@forEach
+            }
+
+            result.append(context.format())
+            result.append("\n")
+
+            createForFields(context)
+
+            nestedNodes.forEach {
+                result.append("\n")
+                result.append(context.format())
+            }
+        }
+
+        return result.toString()
+    }
+
+    private fun createForFields(context: ClassContext) {
+        val filterFields: List<PsiElement> = context.fields.filter {
+            isInProject(it.containingFile?.virtualFile!!, project)
+        }
+
+        filterFields.forEach {
+            val variableContext: VariableContext = VariableContextProvider(false, false, false).from(it)
+            variableContext.text
+            val classContext: ClassContext = ClassContextProvider(false).from(it)
+            if (classContext.name != null) {
+                nestedNodes[classContext.name] = classContext
+                createForFields(classContext)
+            }
+        }
+    }
 }
diff --git a/src/main/kotlin/cc/unitmesh/devti/context/base/LLMCodeContextProvider.kt b/src/main/kotlin/cc/unitmesh/devti/context/base/LLMCodeContextProvider.kt
@@ -3,6 +3,6 @@ package cc.unitmesh.devti.context.base;
 import com.intellij.psi.PsiElement
 
 interface LLMCodeContextProvider<T : PsiElement?> {
-    fun from(psiElement: T): LLMCodeContext
+    fun from(psiElement: T): LLMCodeContext?
 }