feat: <python> add basic annotation name support

phodal · Feb 11, 2020 · ccead4d · ccead4d
1 parent b0515c6
commit ccead4d
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 0 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
@@ -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<CodeProperty> {
@@ -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<CodeAnnotation> {
+        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()
+        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
@@ -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
@@ -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")
+    }
 }