Expose Go post-processing utilities

grafana · Aug 1, 2023 · 131d936 · 131d936
1 parent 5a1a452
commit 131d936
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 69 deletions.
diff --git a/encoding/gocode/gen.go b/encoding/gocode/gen.go
@@ -4,25 +4,18 @@ import (
 	"bytes"
 	"embed"
 	"fmt"
-	"go/parser"
-	"go/token"
-	"path/filepath"
 	"strings"
 	"text/template"
 	"time"
 
 	"cuelang.org/go/cue"
 	"cuelang.org/go/pkg/encoding/yaml"
 	"github.com/dave/dst"
-	"github.com/dave/dst/decorator"
 	"github.com/dave/dst/dstutil"
 	"github.com/getkin/kin-openapi/openapi3"
-	"golang.org/x/tools/imports"
-
 	"github.com/grafana/thema"
 	"github.com/grafana/thema/encoding/openapi"
 	"github.com/grafana/thema/internal/deepmap/oapi-codegen/pkg/codegen"
-	"github.com/grafana/thema/internal/util"
 )
 
 // All the parsed templates in the tmpl subdirectory
@@ -129,11 +122,11 @@ func GenerateTypesOpenAPI(sch thema.Schema, cfg *TypeConfigOpenAPI) ([]byte, err
 		return nil, fmt.Errorf("openapi generation failed: %w", err)
 	}
 
-	return postprocessGoFile(genGoFile{
-		path:     fmt.Sprintf("%s_type_gen.go", sch.Lineage().Name()),
-		appliers: applyFuncs,
-		in:       []byte(gostr),
-		errifadd: !cfg.IgnoreDiscoveredImports,
+	return PostprocessGoFile(GenGoFile{
+		Path:     fmt.Sprintf("%s_type_gen.go", sch.Lineage().Name()),
+		Appliers: applyFuncs,
+		In:       []byte(gostr),
+		ErrIfAdd: !cfg.IgnoreDiscoveredImports,
 	})
 }
 
@@ -349,11 +342,11 @@ func GenerateLineageBinding(lin thema.Lineage, cfg *BindingConfig) ([]byte, erro
 		return nil, fmt.Errorf("error executing binding template: %w", err)
 	}
 
-	return postprocessGoFile(genGoFile{
-		path:     fmt.Sprintf("%s_binding_gen.go", strings.ToLower(lin.Name())),
-		appliers: cfg.ApplyFuncs,
-		in:       buf.Bytes(),
-		errifadd: !cfg.IgnoreDiscoveredImports,
+	return PostprocessGoFile(GenGoFile{
+		Path:     fmt.Sprintf("%s_binding_gen.go", strings.ToLower(lin.Name())),
+		Appliers: cfg.ApplyFuncs,
+		In:       buf.Bytes(),
+		ErrIfAdd: !cfg.IgnoreDiscoveredImports,
 	})
 }
 
@@ -396,55 +389,3 @@ type bindingVars struct {
 
 	TargetSchemaVersion thema.SyntacticVersion
 }
-
-type genGoFile struct {
-	errifadd bool
-	path     string
-	appliers []dstutil.ApplyFunc
-	in       []byte
-}
-
-func postprocessGoFile(cfg genGoFile) ([]byte, error) {
-	fname := util.SanitizeLabelString(filepath.Base(cfg.path))
-	buf := new(bytes.Buffer)
-	fset := token.NewFileSet()
-	gf, err := decorator.ParseFile(fset, fname, string(cfg.in), parser.ParseComments)
-	if err != nil {
-		return nil, fmt.Errorf("error parsing generated file: %w", err)
-	}
-
-	for _, af := range cfg.appliers {
-		dstutil.Apply(gf, af, nil)
-	}
-
-	err = decorator.Fprint(buf, gf)
-	if err != nil {
-		return nil, fmt.Errorf("error formatting generated file: %w", err)
-	}
-
-	byt, err := imports.Process(fname, buf.Bytes(), nil)
-	if err != nil {
-		return nil, fmt.Errorf("goimports processing of generated file failed: %w", err)
-	}
-
-	if cfg.errifadd {
-		// Compare imports before and after; warn about performance if some were added
-		gfa, _ := parser.ParseFile(fset, fname, string(byt), parser.ParseComments)
-		imap := make(map[string]bool)
-		for _, im := range gf.Imports {
-			imap[im.Path.Value] = true
-		}
-		var added []string
-		for _, im := range gfa.Imports {
-			if !imap[im.Path.Value] {
-				added = append(added, im.Path.Value)
-			}
-		}
-
-		if len(added) != 0 {
-			// TODO improve the guidance in this error if/when we better abstract over imports to generate
-			return nil, fmt.Errorf("goimports added the following import statements to %s: \n\t%s\nRelying on goimports to find imports significantly slows down code generation. Either add these imports with an AST manipulation in cfg.ApplyFuncs, or set cfg.IgnoreDiscoveredImports to true", cfg.path, strings.Join(added, "\n\t"))
-		}
-	}
-	return byt, nil
-}
diff --git a/encoding/gocode/postprocessing.go b/encoding/gocode/postprocessing.go
@@ -0,0 +1,68 @@
+package gocode
+
+import (
+	"bytes"
+	"fmt"
+	"go/parser"
+	"go/token"
+	"path/filepath"
+	"strings"
+
+	"github.com/dave/dst/decorator"
+	"github.com/dave/dst/dstutil"
+	"golang.org/x/tools/imports"
+
+	"github.com/grafana/thema/internal/util"
+)
+
+type GenGoFile struct {
+	ErrIfAdd bool
+	Path     string
+	Appliers []dstutil.ApplyFunc
+	In       []byte
+}
+
+func PostprocessGoFile(cfg GenGoFile) ([]byte, error) {
+	fname := util.SanitizeLabelString(filepath.Base(cfg.Path))
+	buf := new(bytes.Buffer)
+	fset := token.NewFileSet()
+	gf, err := decorator.ParseFile(fset, fname, string(cfg.In), parser.ParseComments)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing generated file: %w", err)
+	}
+
+	for _, af := range cfg.Appliers {
+		dstutil.Apply(gf, af, nil)
+	}
+
+	err = decorator.Fprint(buf, gf)
+	if err != nil {
+		return nil, fmt.Errorf("error formatting generated file: %w", err)
+	}
+
+	byt, err := imports.Process(fname, buf.Bytes(), nil)
+	if err != nil {
+		return nil, fmt.Errorf("goimports processing of generated file failed: %w", err)
+	}
+
+	if cfg.ErrIfAdd {
+		// Compare imports before and after; warn about performance if some were added
+		gfa, _ := parser.ParseFile(fset, fname, string(byt), parser.ParseComments)
+		imap := make(map[string]bool)
+		for _, im := range gf.Imports {
+			imap[im.Path.Value] = true
+		}
+		var added []string
+		for _, im := range gfa.Imports {
+			if !imap[im.Path.Value] {
+				added = append(added, im.Path.Value)
+			}
+		}
+
+		if len(added) != 0 {
+			// TODO improve the guidance in this error if/when we better abstract over imports to generate
+			return nil, fmt.Errorf("goimports added the following import statements to %s: \n\t%s\nRelying on goimports to find imports significantly slows down code generation. Either add these imports with an AST manipulation in cfg.ApplyFuncs, or set cfg.IgnoreDiscoveredImports to true", cfg.Path, strings.Join(added, "\n\t"))
+		}
+	}
+	return byt, nil
+}