autogen.go

// Package autogen automatically generates files from templates.
//
// Usage
//
// Create a main package somewhere in your project. In that package, locate templates for generating
// files. These templates should have the generated file names suffixed with ".gotmpl". The
// generated files will be located by default in the parent directory, but this location can be
// customized with the `Location` option. The main package should define the variables to be used
// in the templates and run the `autogen.Execute` function with this variable.
// The main package should include a `//go:generate go run .` comment. All that left to do is to run
// `go generate ./...` from the Go module's root.
//
// See example at the ./example directory.
package autogen

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/posener/script"
)

// Go-generate environment variables.
var (
	Arch    = os.Getenv("GOARCH")
	OS      = os.Getenv("GOOS")
	File    = os.Getenv("GOFILE") // Modified to be file path relative to module path.
	Line    = os.Getenv("GOLINE")
	Package = os.Getenv("GOPACKAGE")
)

var (
	// ModulePath is the absolute path of the current Go module.
	ModulePath string

	// absPath holds the absolute path of the main package that called autogen.
	absPath string
)

type config struct {
	outputPath string
}

// Location sets the output path location. Default is the parent directory of the generating file.
func Location(path string) func(c *config) {
	return func(c *config) {
		c.outputPath = path
	}
}

// Execute the autogeneration process with the given value.
func Execute(v interface{}, options ...func(*config)) error {
	var c config
	for _, o := range options {
		o(&c)
	}

	log.Printf("Autogen of Go generate rule: %s:%s", File, Line)

	if c.outputPath == "" {
		// Take parent directory for output.
		c.outputPath = filepath.Dir(absPath)
	}

	log.Printf("Output location: %s", c.outputPath)

	// Store go files names to eventually format them.
	var goFiles []string

	// Load templates from local path.
	templates, err := template.ParseGlob("*.gotmpl")
	if err != nil {
		return fmt.Errorf("loading templates: %s", err)
	}

	for _, t := range templates.Templates() {

		log.Printf("Processing template: %s", t.Name())

		// Calculate output file path.
		out := t.Name()
		out = out[:strings.LastIndex(out, ".")] // Remove .gotmpl suffix.
		out = filepath.Join(c.outputPath, out)

		log.Printf("Writing: %s", out)

		f, err := os.Create(out)
		if err != nil {
			return fmt.Errorf("creating file %s: %s", out, err)
		}
		defer f.Close()

		ext := filepath.Ext(out)

		// Write autogenerated comment if possible.
		if prefix := commentPrefix(ext); prefix != "" {
			_, err = f.WriteString(fmt.Sprintf("%s Autogenerated by go run %s. DO NOT EDIT. \n\n", prefix, File))
			if err != nil {
				return fmt.Errorf("writing comment string: %s", err)
			}
		}

		// Execute the template.
		err = t.Execute(f, v)
		if err != nil {
			return fmt.Errorf("executing template %s: %s", t.Name(), err)
		}

		if ext == ".go" {
			goFiles = append(goFiles, out)
		}
	}

	// Format Go files.
	if len(goFiles) > 0 {
		log.Printf("Formatting Go files.")
		args := append([]string{"-w"}, goFiles...)
		err := script.ExecHandleStderr(os.Stderr, "goimports", args...).ToStdout()
		if err != nil {
			return fmt.Errorf("formatting go files: %s", err)
		}
	}

	return nil
}

// commentPrefix returns a comment prefix according to a given file type extension.
func commentPrefix(ext string) string {
	switch ext {
	case ".go":
		return "//"
	case ".yml", ".yaml":
		return "#"
	default:
		return ""
	}
}

func init() {
	var err error
	absPath, err = filepath.Abs(".")
	if err != nil {
		panic(err)
	}

	// Lookup module path.
	for path := absPath; ; {
		if _, err := os.Stat(filepath.Join(path, "go.mod")); err == nil {
			ModulePath = path
			break
		}
		if newPath := filepath.Dir(path); newPath == path {
			break
		} else {
			path = newPath
		}
	}

	// Calculate RelPath and File relative to module path.
	relPath, err := filepath.Rel(ModulePath, absPath)
	if err != nil {
		panic(err)
	}
	File = filepath.Join(relPath, File)
}