From 9c063310a099a260a0d708f0e82ca8a617f6ed06 Mon Sep 17 00:00:00 2001
From: Phodal Huang <h@phodal.com>
Date: Tue, 11 Feb 2020 19:49:19 +0800
Subject: [PATCH] feat: <python> add import stmt support

---
 .../ast/pythonast/PythonFullIdentListener.kt  | 24 +++++++++++++++++++
 .../pythonast/PythonFullIdentListenerTest.kt  | 14 +++++++++++
 2 files changed, 38 insertions(+)

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 9bf834a7..4e04c244 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
@@ -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
index 807e8454..0704eeb4 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
@@ -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)
+    }
 }