gazelle.go

// Package stringer is a Gazelle extension that generates go_stringer rules
// from //go:generate stringer directives in Go source files.
package stringer

import (
	"bufio"
	"flag"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/bazelbuild/bazel-gazelle/language"
	"github.com/bazelbuild/bazel-gazelle/rule"
)

// NewLanguage is called by Gazelle to instantiate this extension.
func NewLanguage() language.Language {
	return &lang{}
}

type lang struct {
	language.BaseLang
}

func (*lang) Name() string { return "stringer" }

func (*lang) Kinds() map[string]rule.KindInfo {
	return map[string]rule.KindInfo{
		"go_stringer": {
			MatchAttrs:     []string{"output"},
			NonEmptyAttrs:  map[string]bool{"src": true, "type": true, "mod": true, "output": true},
			MergeableAttrs: map[string]bool{},
		},
	}
}

func (*lang) Loads() []rule.LoadInfo {
	return []rule.LoadInfo{{
		Name:    "//bazel/rules/go_stringer:defs.bzl",
		Symbols: []string{"go_stringer"},
	}}
}

func (*lang) GenerateRules(args language.GenerateArgs) language.GenerateResult {
	mod := findGomod(args.Config.RepoRoot, args.Dir)
	var rules []*rule.Rule

	for _, f := range args.RegularFiles {
		if !strings.HasSuffix(f, ".go") {
			continue
		}
		directives, err := parseFile(filepath.Join(args.Dir, f))
		if err != nil {
			continue
		}
		for _, d := range directives {
			out := d.output
			if out == "" {
				out = strings.ToLower(strings.SplitN(d.typ, ",", 2)[0]) + "_string.go"
			}
			r := rule.NewRule("go_stringer", strings.TrimSuffix(out, ".go"))
			r.SetAttr("src", f)
			r.SetAttr("type", d.typ)
			r.SetAttr("mod", mod)
			r.SetAttr("output", out)
			if d.trimprefix != "" {
				r.SetAttr("trimprefix", d.trimprefix)
			}
			if d.linecomment {
				r.SetAttr("linecomment", true)
			}
			if d.tags != "" {
				r.SetAttr("go_tags", d.tags)
			}
			rules = append(rules, r)
		}
	}
	if len(rules) == 0 {
		return language.GenerateResult{}
	}
	return language.GenerateResult{
		Gen:     rules,
		Imports: make([]interface{}, len(rules)),
	}
}

type directive struct {
	typ        string
	output     string
	trimprefix string
	linecomment bool
	tags       string
}

func parseFile(path string) ([]directive, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var result []directive
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "//go:generate ") {
			continue
		}
		if d, ok := parseDirective(strings.TrimPrefix(line, "//go:generate ")); ok {
			result = append(result, d)
		}
	}
	return result, scanner.Err()
}

func parseDirective(s string) (directive, bool) {
	fields := strings.Fields(s)
	if len(fields) == 0 {
		return directive{}, false
	}
	var args []string
	switch {
	case isStringerCmd(fields[0]):
		args = fields[1:]
	case fields[0] == "go" && len(fields) >= 3 && fields[1] == "run" && isStringerCmd(fields[2]):
		args = fields[3:]
	default:
		return directive{}, false
	}

	fs := flag.NewFlagSet("stringer", flag.ContinueOnError)
	fs.SetOutput(io.Discard)
	typ := fs.String("type", "", "")
	output := fs.String("output", "", "")
	trimprefix := fs.String("trimprefix", "", "")
	linecomment := fs.Bool("linecomment", false, "")
	tags := fs.String("tags", "", "")
	if err := fs.Parse(args); err != nil || *typ == "" {
		return directive{}, false
	}
	return directive{
		typ:         *typ,
		output:      *output,
		trimprefix:  *trimprefix,
		linecomment: *linecomment,
		tags:        *tags,
	}, true
}

func isStringerCmd(s string) bool {
	return s == "stringer" ||
		strings.HasSuffix(s, "/stringer") ||
		strings.HasSuffix(s, "/cmd/stringer")
}

func findGomod(repoRoot, dir string) string {
	for d := dir; ; {
		if _, err := os.Stat(filepath.Join(d, "go.mod")); err == nil {
			rel, err := filepath.Rel(repoRoot, d)
			if err != nil || rel == "." {
				return "//:go.mod"
			}
			return "//" + filepath.ToSlash(rel) + ":go.mod"
		}
		parent := filepath.Dir(d)
		if parent == d {
			return "//:go.mod"
		}
		d = parent
	}
}