feat: <python> add basic local vars support

chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonAstBaseListener.kt

@@ -7,7 +7,7 @@ import org.antlr.v4.runtime.tree.ParseTree
 
 open class PythonAstBaseListener : PythonParserBaseListener() {
     var currentFunction: CodeFunction = CodeFunction()
-    private var localVars = mutableMapOf<String, String>()
+    var localVars = mutableMapOf<String, String>()
 
     fun buildParameters(listCtx: PythonParser.TypedargslistContext?): Array<CodeProperty> {
         var parameters: Array<CodeProperty> = arrayOf()
@@ -103,29 +103,53 @@ open class PythonAstBaseListener : PythonParserBaseListener() {
         val assignPartCtx = exprCtx.assign_part()
         if (assignPartCtx != null) {
             if (assignPartCtx.ASSIGN() != null) {
-//                println(assignPartCtx.getChild(1).text)
+                val rightObj = buildAssignPart(assignPartCtx)
+                localVars[leftPart] = rightObj
             }
         }
     }
 
-    private fun buildTestContext(testContext: PythonParser.TestContext) {
+    private fun buildAssignPart(assignPartCtx: PythonParser.Assign_partContext): String {
+        var returnAtom = ""
+        for (starExprCtx in assignPartCtx.testlist_star_expr()) {
+            val child = starExprCtx.getChild(0)
+            when (child::class.java.simpleName) {
+                "TestlistContext" -> {
+                    val testlistCtx = child as PythonParser.TestlistContext
+                    for (testContext in testlistCtx.test()) {
+                        returnAtom = this.buildTestContext(testContext)
+                    }
+                }
+            }
+        }
+
+        return returnAtom
+    }
+
+    private fun buildTestContext(testContext: PythonParser.TestContext): String {
+        var returnType = ""
         val exprCtx = testContext.getChild(0).getChild(0).getChild(0)
-        val expr = exprCtx::class.java.simpleName
-        when(expr) {
+        val exprType = exprCtx::class.java.simpleName
+        when (exprType) {
             "ExprContext" -> {
                 val exprCtx = exprCtx as PythonParser.ExprContext
                 val exprChild = exprCtx.getChild(0)::class.java.simpleName
                 when (exprChild) {
                     "AtomContext" -> {
-                        buildAtomExpr(exprCtx)
+                        returnType = buildAtomExpr(exprCtx)
                     }
                 }
             }
+            else -> {
+                println("buildTestContext -> " + exprType)
+            }
         }
+
+        return returnType
     }
 
-    private fun buildAtomExpr(exprCtx: PythonParser.ExprContext) {
-        var funcCalls : Array<CodeCall> = arrayOf()
+    private fun buildAtomExpr(exprCtx: PythonParser.ExprContext): String {
+        var funcCalls: Array<CodeCall> = arrayOf()
         val atomName = exprCtx.atom().text
         for (trailerContext in exprCtx.trailer()) {
             if (trailerContext.DOT() != null) {
@@ -146,5 +170,7 @@ open class PythonAstBaseListener : PythonParserBaseListener() {
         }
 
         currentFunction.FunctionCalls = funcCalls
+
+        return atomName
     }
 }

chapi-ast-python/src/main/kotlin/chapi/ast/pythonast/PythonFullIdentListener.kt

@@ -102,6 +102,7 @@ class PythonFullIdentListener(var fileName: String) : PythonAstBaseListener() {
     }
 
     override fun exitFuncdef(ctx: PythonParser.FuncdefContext?) {
+        currentFunction.addVarsFromMap(this.localVars)
         if (currentNode.NodeName == "") {
             defaultNode.Functions += currentFunction
         } else {

chapi-ast-python/src/test/kotlin/chapi/ast/pythonast/PythonAnalyserTest.kt

@@ -32,4 +32,19 @@ async def show(str):
         assertEquals(firstFunc.FunctionCalls[0].FunctionName, "print")
         assertEquals(firstFunc.FunctionCalls[0].NodeName, "")
     }
+
+    @Test
+    internal fun shouldIdentLocalVarsForFuncCall() {
+        val code = """
+def build(): 
+    p = Person()
+    p.setName("Hunter").setAge(24).setSSN("111-22-3333")
+"""
+
+        val codeFile = PythonAnalyser().analysis(code, "")
+        val firstFunc = codeFile.DataStructures[0].Functions[0]
+        assertEquals(firstFunc.LocalVariables.size, 1)
+        assertEquals(firstFunc.LocalVariables[0].TypeValue, "p")
+        assertEquals(firstFunc.LocalVariables[0].TypeType, "Person")
+    }
 }