feat(context): add support for annotations in ClassContext

This commit adds support for annotations in the ClassContext class. The KotlinClassContextBuilder now retrieves the annotation entries from the PsiElement and maps them to a list of strings. The ClassContext constructor is updated to include the annotations parameter. The ClassContext's format() method is also updated to include the annotations in the output. A new test case is added to verify the functionality.
unit-mesh · Jan 10, 2024 · abc5305 · abc5305
1 parent c5903df
commit abc5305
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 4 deletions.
diff --git a/kotlin/src/main/kotlin/cc/unitmesh/kotlin/context/KotlinClassContextBuilder.kt b/kotlin/src/main/kotlin/cc/unitmesh/kotlin/context/KotlinClassContextBuilder.kt
@@ -28,8 +28,12 @@ class KotlinClassContextBuilder : ClassContextBuilder {
         val usages =
             if (gatherUsages) JavaContextCollection.findUsages(psiElement as PsiNameIdentifierOwner) else emptyList()
 
+        val annotations: List<String> = psiElement.annotationEntries.mapNotNull {
+            it.text
+        }
+
         val displayName = psiElement.fqName?.asString() ?: psiElement.name ?: ""
-        return ClassContext(psiElement, text, name, ktNamedFunctions, allFields, null, usages, displayName = displayName)
+        return ClassContext(psiElement, text, name, ktNamedFunctions, allFields, null, usages, displayName = displayName, annotations)
     }
 
     companion object {

diff --git a/kotlin/src/test/kotlin/cc/unitmesh/kotlin/context/KotlinClassContextBuilderTest.kt b/kotlin/src/test/kotlin/cc/unitmesh/kotlin/context/KotlinClassContextBuilderTest.kt
@@ -0,0 +1,57 @@
+package cc.unitmesh.kotlin.context;
+
+import cc.unitmesh.devti.context.ClassContext
+import cc.unitmesh.devti.context.builder.ClassContextBuilder
+import cc.unitmesh.idea.context.JavaContextCollection
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiNameIdentifierOwner
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.testFramework.LightPlatformTestCase
+import junit.framework.TestCase
+import org.jetbrains.annotations.Nullable
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtParameter
+import org.jetbrains.kotlin.psi.KtPsiFactory
+
+class KotlinClassContextBuilderTest: LightPlatformTestCase() {
+
+    fun testShould_return_functions_from_kt_class_or_object() {
+        val code = """
+            package cc.unitmesh.untitled.demo.controller
+
+            import org.springframework.web.bind.annotation.*
+            
+            @RestController
+            @RequestMapping("/user")
+            class UserController() {
+                @GetMapping
+                fun getUsers(): UserDTO {
+                    return UserDTO(1L, "username", "email", "name", "surname")
+                }
+            }
+            
+            data class UserDTO(
+                val id: Long? = null,
+                val username: String,
+                val email: String,
+                val name: String,
+                val surname: String? = null,
+            )
+            """.trimIndent()
+
+        val createFile = KtPsiFactory(project).createFile("UserController.kt", code)
+        val controller = PsiTreeUtil.findChildOfType(createFile, KtClassOrObject::class.java)!!
+        // given
+        val builder = KotlinClassContextBuilder()
+
+        // then
+        val classContext = builder.getClassContext(controller, false)!!
+        TestCase.assertEquals(classContext.format(), "'package: cc.unitmesh.untitled.demo.controller.UserController\n" +
+                "'@RestController, @RequestMapping(\"/user\")\n" +
+                "class UserController {\n" +
+                "  \n" +
+                "  \n" +
+                "}")
+    }
+}
diff --git a/kotlin/src/test/kotlin/cc/unitmesh/kotlin/provider/KotlinTestDataBuilderTest.kt b/kotlin/src/test/kotlin/cc/unitmesh/kotlin/provider/KotlinTestDataBuilderTest.kt
@@ -12,7 +12,7 @@ class KotlinTestDataBuilderTest : LightPlatformTestCase() {
     }
 
     // test will fail if 222
-    fun shouldReturnLangFileSuffix() {
+    fun testShouldReturnLangFileSuffix() {
         val code = """
             package cc.unitmesh.untitled.demo.controller
 

diff --git a/src/main/kotlin/cc/unitmesh/devti/context/ClassContext.kt b/src/main/kotlin/cc/unitmesh/devti/context/ClassContext.kt
@@ -13,7 +13,8 @@ class ClassContext(
     val fields: List<PsiElement> = emptyList(),
     val superClasses: List<String>? = null,
     val usages: List<PsiReference> = emptyList(),
-    val displayName: String? = null
+    val displayName: String? = null,
+    val annotations: List<String> = mutableListOf(),
 ) : NamedElementContext(root, text, name) {
     private fun getFieldNames(): List<String> = fields.mapNotNull {
         VariableContextProvider(false, false, false).from(it).name
@@ -36,9 +37,14 @@ class ClassContext(
             }
 
         val filePath = displayName ?: runReadAction { root.containingFile?.virtualFile?.path }
+        val annotations = if (annotations.isEmpty()) {
+            ""
+        } else {
+            "\n'" + annotations.joinToString(separator = ", ")
+        }
 
         return """
-        |'package: $filePath
+        |'package: $filePath$annotations
         |class $className$superClasses {
         |  $classFields
         |  $methodSignatures