From 937a95068fdf814a1e5cda2cb9971712d821730c Mon Sep 17 00:00:00 2001
From: Phodal Huang <h@phodal.com>
Date: Tue, 11 Feb 2020 17:00:48 +0800
Subject: [PATCH] feat: <python> add annotation key value support

---
 .../ast/pythonast/PythonAstBaseListener.kt    | 28 +++++++++++++++++--
 .../pythonast/PythonFullIdentListenerTest.kt  | 15 ++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonAstBaseListener.kt b/chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonAstBaseListener.kt
index dd639bfa..9a4770f8 100644
--- a/chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonAstBaseListener.kt
+++ b/chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonAstBaseListener.kt
@@ -2,6 +2,7 @@ package chapi.ast.pythonast
 
 import chapi.ast.antlr.PythonParser
 import chapi.ast.antlr.PythonParserBaseListener
+import domain.core.AnnotationKeyValue
 import domain.core.CodeAnnotation
 import domain.core.CodeProperty
 import org.antlr.v4.runtime.tree.ParseTree
@@ -43,12 +44,12 @@ open class PythonAstBaseListener : PythonParserBaseListener() {
     }
 
     fun buildAnnotationsByIndex(ctx: PythonParser.ClassdefContext, ctxIndex: Int): Array<CodeAnnotation> {
-        var nodes : Array<PythonParser.DecoratorContext> = arrayOf()
+        var nodes: Array<PythonParser.DecoratorContext> = arrayOf()
         for (i in 0 until ctxIndex) {
             nodes += ctx.parent.getChild(i) as PythonParser.DecoratorContext
         }
 
-        var annotations : Array<CodeAnnotation> = arrayOf()
+        var annotations: Array<CodeAnnotation> = arrayOf()
         for (node in nodes) {
             annotations += this.buildAnnotation(node)
         }
@@ -61,6 +62,29 @@ open class PythonAstBaseListener : PythonParserBaseListener() {
             Name = node.dotted_name().text
         )
 
+        if (node.arglist() != null) {
+            codeAnnotation.KeyValues = this.buildArgList(node.arglist())
+        }
+
         return codeAnnotation
     }
+
+    private fun buildArgList(arglistCtx: PythonParser.ArglistContext?): Array<AnnotationKeyValue> {
+        var arguments: Array<AnnotationKeyValue> = arrayOf()
+        for (argCtx in arglistCtx!!.argument()) {
+            val key = argCtx.test(0).text
+            var value = ""
+            if (argCtx.test().size > 1) {
+                value = argCtx.test(1).text
+            }
+
+            val annotation = AnnotationKeyValue(
+                Key = key,
+                Value = value
+            )
+            arguments += annotation
+        }
+
+        return arguments
+    }
 }
diff --git a/chapi-ast-python/src/test/kotlin/chapi/ast/pythonast/PythonFullIdentListenerTest.kt b/chapi-ast-python/src/test/kotlin/chapi/ast/pythonast/PythonFullIdentListenerTest.kt
index 99e323ba..0767d72d 100644
--- a/chapi-ast-python/src/test/kotlin/chapi/ast/pythonast/PythonFullIdentListenerTest.kt
+++ b/chapi-ast-python/src/test/kotlin/chapi/ast/pythonast/PythonFullIdentListenerTest.kt
@@ -135,4 +135,19 @@ class foo:
         assertEquals(codeFile.DataStructures[0].Annotations[0].Name, "cache")
         assertEquals(codeFile.DataStructures[0].Annotations[1].Name, "decorator")
     }
+
+    @Test
+    internal fun shouldIdentifyMultipleClassAnnotationKeyValue() {
+        val code = """
+@cache(key="value")
+class multiple_annotation():
+    pass
+
+"""
+
+        val codeFile = PythonAnalyser().analysis(code, "")
+        assertEquals(codeFile.DataStructures[0].Annotations[0].Name, "cache")
+        assertEquals(codeFile.DataStructures[0].Annotations[0].KeyValues[0].Key, "key")
+        assertEquals(codeFile.DataStructures[0].Annotations[0].KeyValues[0].Value, "\"value\"")
+    }
 }