feat: <go> add function multiple returnType support

phodal · Feb 12, 2020 · 5a03bb4 · 5a03bb4
1 parent 9853b2e
commit 5a03bb4
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 7 deletions.
diff --git a/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoAstListener.kt b/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoAstListener.kt
@@ -23,16 +23,27 @@ open class GoAstListener : GoParserBaseListener() {
         return ""
     }
 
-    protected fun buildReturnTypeFromSignature(codeFunction: CodeFunction, signatureContext: GoParser.SignatureContext?): Array<CodeProperty> {
+    protected fun buildReturnTypeFromSignature(
+        codeFunction: CodeFunction,
+        signatureContext: GoParser.SignatureContext?
+    ): Array<CodeProperty> {
         val result = signatureContext!!.result()
-        var returns : Array<CodeProperty> = arrayOf()
+        var returns: Array<CodeProperty> = arrayOf()
         if (result != null) {
-            val returnType = CodeProperty(
-                TypeType = result.text,
-                TypeValue = ""
-            )
+            val parameters = result.parameters()
+            if (parameters != null) {
+                for (parameterDeclContext in parameters.parameterDecl()) {
+                    val typeType = parameterDeclContext.text
+                    val returnType = CodeProperty(TypeType = typeType, TypeValue = "")
+                    returns += returnType
+                }
+            }
 
-            returns += returnType
+            if (result.type_() != null) {
+                val typeType = result.type_().text
+                val returnType = CodeProperty(TypeType = typeType, TypeValue = "")
+                returns += returnType
+            }
         }
 
         return returns;

diff --git a/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt b/chapi-ast-go/src/main/kotlin/chapi/ast/goast/GoFullIdentListener.kt
@@ -45,6 +45,8 @@ class GoFullIdentListener(var fileName: String) : GoAstListener() {
             Package = codeFile.PackageName
         )
 
+        codeFunction.MultipleReturns = this.buildReturnTypeFromSignature(codeFunction, ctx.signature())
+
         defaultNode.Functions += codeFunction
     }
 

diff --git a/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt b/chapi-ast-go/src/test/kotlin/chapi/ast/goast/GoFullIdentListenerTest.kt
@@ -142,4 +142,19 @@ func add(x int, y int) int {
         assertEquals(codeFile.DataStructures[0].Functions.size, 1)
         assertEquals(codeFile.DataStructures[0].Functions[0].Name, "add")
     }
+
+    @Test
+    internal fun shouldIdentifyFunctionMultipleReturnType() {
+        var code = """
+package main
+
+func get(x int, y int) (int, int) {
+	return x, y
+}
+"""
+        val codeFile = GoAnalyser().analysis(code, "")
+        assertEquals(codeFile.DataStructures.size, 1)
+        assertEquals(codeFile.DataStructures[0].Functions.size, 1)
+        assertEquals(codeFile.DataStructures[0].Functions[0].MultipleReturns.size, 2)
+    }
 }
-Original file line number
+Diff line change
@@ Expand Up @@
                 Package = codeFile.PackageName
             )
+            codeFunction.MultipleReturns = this.buildReturnTypeFromSignature(codeFunction, ctx.signature())
             defaultNode.Functions += codeFunction
         }
@@ Expand Down @@