feat: <python> add basic inheritance

phodal · Feb 11, 2020 · 268cfe7 · 268cfe7
1 parent 5be1b82
commit 268cfe7
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
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
@@ -15,6 +15,12 @@ class PythonFullIdentListener(var fileName: String) : PythonAstBaseListener() {
         currentNode = CodeDataStruct(
             NodeName = ctx!!.name().text
         )
+
+        if (ctx.arglist() != null) {
+            for (argumentContext in ctx.arglist().argument()) {
+                currentNode.MultipleExtend += argumentContext.text
+            }
+        }
     }
 
     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
@@ -25,4 +25,19 @@ internal class PythonFullIdentListenerTest {
         assertEquals(codeFile.DataStructures.size, 1)
         assertEquals(codeFile.DataStructures[0].NodeName, "foo")
     }
+
+    @Test
+    internal fun shouldIdentClassInheritance() {
+        val code = """
+class Student(Person):
+    pass
+        """
+        val codeFile = PythonAnalyser().analysis(code, "")
+
+        assertEquals(codeFile.DataStructures.size, 1)
+        assertEquals(codeFile.DataStructures[0].NodeName, "Student")
+        assertEquals(codeFile.DataStructures[0].MultipleExtend.size, 1)
+        assertEquals(codeFile.DataStructures[0].Extend, "")
+        assertEquals(codeFile.DataStructures[0].MultipleExtend[0], "Person")
+    }
 }