feat: <python> add import as support

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

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
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
@@ -27,7 +27,7 @@ class PythonFullIdentListener(var fileName: String) : PythonAstBaseListener() {
             codeImport.UsageName += firstNameCtx.name().text
         }
 
-        for (i in 1 until dotNames.size - 1) {
+        for (i in 1 until dotNames.size) {
             codeImport.UsageName += dotNames[i].text
         }
 

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
@@ -186,4 +186,19 @@ import collections.abc
         assertEquals(codeFile.Imports[0].Source, "collections.abc")
         assertEquals(codeFile.Imports[0].UsageName.size, 0)
     }
+    @Test
+    internal fun shouldIdentImportAs() {
+        val code = """
+import collections.abc as ss, itertools
+
+"""
+
+        val codeFile = PythonAnalyser().analysis(code, "")
+
+        assertEquals(codeFile.Imports.size, 1)
+        assertEquals(codeFile.Imports[0].Source, "collections.abc")
+        assertEquals(codeFile.Imports[0].UsageName.size, 2)
+        assertEquals(codeFile.Imports[0].UsageName[0], "ss")
+        assertEquals(codeFile.Imports[0].UsageName[1], "itertools")
+    }
 }