feat: <python> add func annotation support

phodal · Feb 11, 2020 · 74b9394 · 74b9394
1 parent 937a950
commit 74b9394
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 1 deletion.
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
@@ -43,7 +43,7 @@ open class PythonAstBaseListener : PythonParserBaseListener() {
         return 0
     }
 
-    fun buildAnnotationsByIndex(ctx: PythonParser.ClassdefContext, ctxIndex: Int): Array<CodeAnnotation> {
+    fun buildAnnotationsByIndex(ctx: ParseTree, ctxIndex: Int): Array<CodeAnnotation> {
         var nodes: Array<PythonParser.DecoratorContext> = arrayOf()
         for (i in 0 until ctxIndex) {
             nodes += ctx.parent.getChild(i) as PythonParser.DecoratorContext

diff --git a/chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonFullIdentListener.kt b/chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonFullIdentListener.kt
@@ -49,6 +49,11 @@ class PythonFullIdentListener(var fileName: String) : PythonAstBaseListener() {
             currentFunction.Modifiers += ctx.ASYNC().text
         }
 
+        val ctxIndex = this.getNodeIndex(ctx)
+        if (ctxIndex > 0) {
+            currentFunction.Annotations = this.buildAnnotationsByIndex(ctx, ctxIndex)
+        }
+
         if (ctx.typedargslist() != null) {
             currentFunction.Parameters = this.buildParameters(ctx.typedargslist())
         }

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
@@ -150,4 +150,26 @@ class multiple_annotation():
         assertEquals(codeFile.DataStructures[0].Annotations[0].KeyValues[0].Key, "key")
         assertEquals(codeFile.DataStructures[0].Annotations[0].KeyValues[0].Value, "\"value\"")
     }
+
+    @Test
+    internal fun shouldIdentifyMultipleFunctionAnnotationKeyValue() {
+        val code = """
+@accepts(int, int)
+@returns(float)
+def bar(low,high):
+    pass
+
+"""
+
+        val codeFile = PythonAnalyser().analysis(code, "")
+        val firstFunc = codeFile.DataStructures[0].Functions[0]
+
+        assertEquals(firstFunc.Annotations[0].Name, "accepts")
+        assertEquals(firstFunc.Annotations[0].KeyValues.size, 2)
+        assertEquals(firstFunc.Annotations[0].KeyValues[0].Key, "int")
+        assertEquals(firstFunc.Annotations[0].KeyValues[1].Key, "int")
+
+        assertEquals(firstFunc.Annotations[1].Name, "returns")
+        assertEquals(firstFunc.Annotations[1].KeyValues[0].Key, "float")
+    }
 }