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 9a4770f8..e24adfab 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 @@ -43,7 +43,7 @@ open class PythonAstBaseListener : PythonParserBaseListener() { return 0 } - fun buildAnnotationsByIndex(ctx: PythonParser.ClassdefContext, ctxIndex: Int): Array { + fun buildAnnotationsByIndex(ctx: ParseTree, ctxIndex: Int): Array { var nodes: Array = 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 index 474c4a37..9bf834a7 100644 --- 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 index 0767d72d..807e8454 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 @@ -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") + } }