Merge pull request #48 from csueiras/fix-47-multi-return-fn-arg

Fixes #47: Multi-return function arguments aren't properly supported in codegen
csueiras · Mar 29, 2021 · 589dfbb · 589dfbb
2 parents 0f5e850 + 5d5e84c
commit 589dfbb
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 1 deletion.
diff --git a/internal/generator/method/method.go b/internal/generator/method/method.go
@@ -220,7 +220,13 @@ func toType(t types.Type, variadic bool) (jen.Code, error) {
 			}
 			returnTypes = append(returnTypes, tt)
 		}
-		return jen.Func().Params(paramTypes...).Add(returnTypes...), nil
+		if len(returnTypes) == 0 {
+			return jen.Func().Params(paramTypes...), nil
+		}
+		if len(returnTypes) > 1 {
+			return jen.Func().Params(paramTypes...).Parens(jen.List(returnTypes...)), nil
+		}
+		return jen.Func().Params(paramTypes...).Add(returnTypes[0]), nil
 	default:
 		return nil, fmt.Errorf("type not handled: %T", v)
 	}

diff --git a/internal/generator/method/method_test.go b/internal/generator/method/method_test.go
@@ -117,6 +117,28 @@ func TestNewMethod(t *testing.T) {
 				ReturnTypes:           []jen.Code{jen.Id("error")},
 			},
 		},
+		{
+			name: "Fn(arg0 func() (string, error))",
+			args: args{
+				name: "Fn",
+				signature: types.NewSignature(nil, types.NewTuple(
+					types.NewVar(token.NoPos, nil, "myArg", types.NewSignature(nil, types.NewTuple(), types.NewTuple(
+						types.NewVar(token.NoPos, nil, "", types.Typ[types.String]),
+						types.NewVar(token.NoPos, nil, "", rtypes.ErrType),
+					), false)),
+				), types.NewTuple(), false),
+			},
+			want: &method.Method{
+				Name:           "Fn",
+				HasContext:     false,
+				ReturnsError:   false,
+				ParameterNames: []string{"arg0"},
+				ParametersNameAndType: []jen.Code{
+					jen.Id("arg0").Add(jen.Func().Params().Parens(jen.List(jen.Id("string"), jen.Id("error")))),
+				},
+				ReturnTypes: []jen.Code{},
+			},
+		},
 		{
 			name: "Fn(ctx context.Context, arg1 string) (string, error)",
 			args: args{
@@ -214,6 +236,7 @@ func TestNewMethod(t *testing.T) {
 			if tt.want.ReturnErrorIndex != nil {
 				require.Equal(t, *tt.want.ReturnErrorIndex, *got.ReturnErrorIndex)
 			}
+
 			require.ElementsMatch(t, tt.want.ParameterNames, got.ParameterNames)
 			require.ElementsMatch(t, tt.want.ParametersNameAndType, got.ParametersNameAndType)
 			require.ElementsMatch(t, tt.want.ReturnTypes, got.ReturnTypes)