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 0251a274..dd639bfa 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,7 +2,9 @@ package chapi.ast.pythonast import chapi.ast.antlr.PythonParser import chapi.ast.antlr.PythonParserBaseListener +import domain.core.CodeAnnotation import domain.core.CodeProperty +import org.antlr.v4.runtime.tree.ParseTree open class PythonAstBaseListener : PythonParserBaseListener() { fun buildParameters(listCtx: PythonParser.TypedargslistContext?): Array { @@ -25,4 +27,40 @@ open class PythonAstBaseListener : PythonParserBaseListener() { return parameters } + + fun getNodeIndex(node: ParseTree?): Int { + if (node == null || node.parent == null) { + return -1 + } + + val parent = node.parent + for (i in 0 until parent.childCount) { + if (parent.getChild(i) == node) { + return i + } + } + return 0 + } + + fun buildAnnotationsByIndex(ctx: PythonParser.ClassdefContext, ctxIndex: Int): Array { + var nodes : Array = arrayOf() + for (i in 0 until ctxIndex) { + nodes += ctx.parent.getChild(i) as PythonParser.DecoratorContext + } + + var annotations : Array = arrayOf() + for (node in nodes) { + annotations += this.buildAnnotation(node) + } + + return annotations + } + + fun buildAnnotation(node: PythonParser.DecoratorContext): CodeAnnotation { + val codeAnnotation = CodeAnnotation( + Name = node.dotted_name().text + ) + + return codeAnnotation + } } 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 d3578c2b..474c4a37 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 @@ -26,6 +26,11 @@ class PythonFullIdentListener(var fileName: String) : PythonAstBaseListener() { currentNode.MultipleExtend += argumentContext.text } } + + val ctxIndex = this.getNodeIndex(ctx) + if (ctxIndex > 0) { + currentNode.Annotations = this.buildAnnotationsByIndex(ctx, ctxIndex) + } } override fun exitClassdef(ctx: PythonParser.ClassdefContext?) { 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 783a850a..55537f0b 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 @@ -103,4 +103,18 @@ def printinfo( name, age = 35): assertEquals(codeFile.DataStructures[0].Functions[0].Parameters[1].TypeValue, "age") assertEquals(codeFile.DataStructures[0].Functions[0].Parameters[1].DefaultValue, "35") } + @Test + internal fun shouldIdentifyClassAnnotation() { + val code = """ +@decorator +class foo: + pass + +""" + + val codeFile = PythonAnalyser().analysis(code, "") + assertEquals(codeFile.DataStructures[0].NodeName, "foo") + assertEquals(codeFile.DataStructures[0].Annotations.size, 1) + assertEquals(codeFile.DataStructures[0].Annotations[0].Name, "decorator") + } }