Collect imports

lokerpc/codegen/go.go

@@ -1,7 +1,7 @@
 package codegen
 
 import (
-	"bufio"
+	"bytes"
 	"fmt"
 	"io"
 	"regexp"
@@ -11,12 +11,12 @@ import (
 	jtd "github.com/jsontypedef/json-typedef-go"
 )
 
-func GenGoType(schema jtd.Schema) string {
+func GenGoType(schema jtd.Schema, imports map[string]struct{}) string {
 	var t string
 
 	for k, v := range schema.Definitions {
 		t += "\n"
-		t += "type " + goFieldName(k) + " " + GenGoType(v) + "\n"
+		t += "type " + goFieldName(k) + " " + GenGoType(v, imports) + "\n"
 	}
 
 	switch schema.Form() {
@@ -28,6 +28,7 @@ func GenGoType(schema jtd.Schema) string {
 			t += "string"
 		case jtd.TypeTimestamp:
 			t += "time.Time"
+			imports["time"] = struct{}{}
 		case jtd.TypeInt8:
 			t += "int8"
 		case jtd.TypeInt16:
@@ -48,16 +49,16 @@ func GenGoType(schema jtd.Schema) string {
 			t += "bool"
 		}
 	case jtd.FormElements:
-		t += "[]" + GenGoType(*schema.Elements)
+		t += "[]" + GenGoType(*schema.Elements, imports)
 	case jtd.FormValues:
-		t += "map[string]" + GenGoType(*schema.Values)
+		t += "map[string]" + GenGoType(*schema.Values, imports)
 	case jtd.FormProperties:
 		t += "struct {\n"
 		for _, k := range sortedKeys(schema.Properties) {
-			t += "\t" + goFieldName(k) + " " + GenGoType(schema.Properties[k]) + "`json:\"" + k + "\"`\n"
+			t += "\t" + goFieldName(k) + " " + GenGoType(schema.Properties[k], imports) + "`json:\"" + k + "\"`\n"
 		}
 		for _, k := range sortedKeys(schema.OptionalProperties) {
-			t += "\t" + goFieldName(k) + " " + GenGoType(schema.OptionalProperties[k]) + "`json:\"" + k + ",omitempty\"`\n"
+			t += "\t" + goFieldName(k) + " " + GenGoType(schema.OptionalProperties[k], imports) + "`json:\"" + k + ",omitempty\"`\n"
 		}
 		t += "}"
 	case jtd.FormDiscriminator:
@@ -80,19 +81,15 @@ func GenGoType(schema jtd.Schema) string {
 func GenGoClient(w io.Writer, meta lokerpc.Meta) error {
 	defOrder := normalise(&meta)
 
-	b := bufio.NewWriter(w)
+	imports := map[string]struct{}{
+		"context": {},
+	}
 
-	fmt.Fprintf(b, "package %s\n", strings.ToLower(strings.ReplaceAll(meta.ServiceName, "-", "")))
-	b.WriteString("\n")
-	b.WriteString("import (\n")
-	b.WriteString("\t\"context\"\n")
-	b.WriteString("\n")
-	b.WriteString("\t\"github.com/LOKE/pkg/lokerpc\"\n")
-	b.WriteString(")\n")
+	var b bytes.Buffer
 
 	for _, k := range defOrder {
 		b.WriteString("\n")
-		fmt.Fprintf(b, "type %s %s;\n", goFieldName(k), GenGoType(meta.Definitions[k]))
+		fmt.Fprintf(&b, "type %s %s;\n", goFieldName(k), GenGoType(meta.Definitions[k], imports))
 	}
 
 	// Service interface
@@ -102,15 +99,15 @@ func GenGoClient(w io.Writer, meta lokerpc.Meta) error {
 	for _, v := range meta.Interfaces {
 		reqType := "any"
 		if v.RequestTypeDef != nil {
-			reqType = GenGoType(*v.RequestTypeDef)
+			reqType = GenGoType(*v.RequestTypeDef, imports)
 		}
 
 		resType := "any"
 		if v.ResponseTypeDef != nil {
 			if v.ResponseTypeDef.Metadata["void"] == true {
 				resType = "struct{}"
 			} else {
-				resType = GenGoType(*v.ResponseTypeDef)
+				resType = GenGoType(*v.ResponseTypeDef, imports)
 
 				if !strings.HasPrefix(resType, "[]") && !strings.HasPrefix(resType, "map[") && !strings.HasPrefix(resType, "*") {
 					resType = "*" + resType
@@ -119,7 +116,7 @@ func GenGoClient(w io.Writer, meta lokerpc.Meta) error {
 		}
 
 		// goDocComment(b, v.Help, "\t")
-		fmt.Fprintf(b, "\t%s(context.Context, %s) (%s, error)\n", goFieldName(v.MethodName), reqType, resType)
+		fmt.Fprintf(&b, "\t%s(context.Context, %s) (%s, error)\n", goFieldName(v.MethodName), reqType, resType)
 	}
 	b.WriteString("}\n")
 
@@ -130,15 +127,15 @@ func GenGoClient(w io.Writer, meta lokerpc.Meta) error {
 	for _, v := range meta.Interfaces {
 		reqType := "any"
 		if v.RequestTypeDef != nil {
-			reqType = GenGoType(*v.RequestTypeDef)
+			reqType = GenGoType(*v.RequestTypeDef, imports)
 		}
 
 		resType := "any"
 		if v.ResponseTypeDef != nil {
 			if v.ResponseTypeDef.Metadata["void"] == true {
 				resType = "struct{}"
 			} else {
-				resType = GenGoType(*v.ResponseTypeDef)
+				resType = GenGoType(*v.ResponseTypeDef, imports)
 
 				if !strings.HasPrefix(resType, "[]") && !strings.HasPrefix(resType, "map[") && !strings.HasPrefix(resType, "*") {
 					resType = "*" + resType
@@ -152,23 +149,35 @@ func GenGoClient(w io.Writer, meta lokerpc.Meta) error {
 		}
 
 		// goDocComment(b, v.Help, "\t")
-		fmt.Fprintf(b, "func (c %sRPCClient) %s(ctx context.Context, req %s) (%s, error) {\n", goFieldName(meta.ServiceName), goFieldName(v.MethodName), reqType, resType)
-		fmt.Fprintf(b, "\tvar res %s\n", varType)
-		fmt.Fprintf(b, "\terr := c.DoRequest(ctx, \"%s\", req, &res)\n", v.MethodName)
-		fmt.Fprintf(b, "\tif err != nil {\n")
-		fmt.Fprintf(b, "\t\treturn nil, err\n")
-		fmt.Fprintf(b, "\t}\n")
+		fmt.Fprintf(&b, "func (c %sRPCClient) %s(ctx context.Context, req %s) (%s, error) {\n", goFieldName(meta.ServiceName), goFieldName(v.MethodName), reqType, resType)
+		fmt.Fprintf(&b, "\tvar res %s\n", varType)
+		fmt.Fprintf(&b, "\terr := c.DoRequest(ctx, \"%s\", req, &res)\n", v.MethodName)
+		fmt.Fprintf(&b, "\tif err != nil {\n")
+		fmt.Fprintf(&b, "\t\treturn nil, err\n")
+		fmt.Fprintf(&b, "\t}\n")
 		if resType == "any" {
-			fmt.Fprintf(b, "\treturn res, nil\n")
+			fmt.Fprintf(&b, "\treturn res, nil\n")
 		} else if strings.HasPrefix(resType, "*") {
-			fmt.Fprintf(b, "\treturn &res, nil\n")
+			fmt.Fprintf(&b, "\treturn &res, nil\n")
 		} else {
-			fmt.Fprintf(b, "\treturn res, nil\n")
+			fmt.Fprintf(&b, "\treturn res, nil\n")
 		}
-		fmt.Fprintf(b, "}\n")
+		fmt.Fprintf(&b, "}\n")
 	}
 
-	return b.Flush()
+	// Write header
+	fmt.Fprintf(w, "package %s\n", strings.ToLower(strings.ReplaceAll(meta.ServiceName, "-", "")))
+	fmt.Fprintf(w, "\nimport (\n")
+
+	for _, im := range sortedKeys(imports) {
+		fmt.Fprintf(w, "\t\"%s\"\n", im)
+	}
+	fmt.Fprintf(w, "\n\t\"github.com/LOKE/pkg/lokerpc\"\n")
+	fmt.Fprintf(w, ")\n\n")
+
+	_, err := io.Copy(w, &b)
+
+	return err
 }
 
 // Regexp that matches word boundaries,

lokerpc/codegen/testdata/nested.json.go

@@ -2,6 +2,7 @@ package nested
 
 import (
 	"context"
+	"time"
 
 	"github.com/LOKE/pkg/lokerpc"
 )