feat: <python> add import stmt support

phodal · Feb 11, 2020 · 9c06331 · 9c06331
1 parent 74b9394
commit 9c06331
Show file tree

Hide file tree

Showing 2 changed files with 38 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
@@ -4,6 +4,7 @@ import chapi.ast.antlr.PythonParser
 import domain.core.CodeDataStruct
 import domain.core.CodeFile
 import domain.core.CodeFunction
+import domain.core.CodeImport
 
 class PythonFullIdentListener(var fileName: String) : PythonAstBaseListener() {
     private var currentFunction: CodeFunction = CodeFunction()
@@ -15,6 +16,29 @@ class PythonFullIdentListener(var fileName: String) : PythonAstBaseListener() {
         NodeName = "default"
     )
 
+    override fun enterImport_stmt(ctx: PythonParser.Import_stmtContext?) {
+        val dotNames = ctx!!.dotted_as_names().dotted_as_name()
+        val firstNameCtx = dotNames[0]
+
+        var codeImport = CodeImport(
+            Source = firstNameCtx.dotted_name().text
+        )
+        if (firstNameCtx.name() != null) {
+            codeImport.UsageName += firstNameCtx.name().text
+        }
+
+        for (i in 1 until dotNames.size - 1) {
+            codeImport.UsageName += dotNames[i].text
+        }
+
+        codeFile.Imports += codeImport
+    }
+
+    override fun enterFrom_stmt(ctx: PythonParser.From_stmtContext?) {
+        super.enterFrom_stmt(ctx)
+        println("super.enterFrom_stmt(ctx)")
+    }
+
     override fun enterClassdef(ctx: PythonParser.ClassdefContext?) {
         hasEnterClass = true
         currentNode = CodeDataStruct(

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
@@ -172,4 +172,18 @@ def bar(low,high):
         assertEquals(firstFunc.Annotations[1].Name, "returns")
         assertEquals(firstFunc.Annotations[1].KeyValues[0].Key, "float")
     }
+
+    @Test
+    internal fun shouldIdentImportDottedName() {
+        val code = """
+import collections.abc
+
+"""
+
+        val codeFile = PythonAnalyser().analysis(code, "")
+
+        assertEquals(codeFile.Imports.size, 1)
+        assertEquals(codeFile.Imports[0].Source, "collections.abc")
+        assertEquals(codeFile.Imports[0].UsageName.size, 0)
+    }
 }