Merge pull request #90 from terrastruct/alixander/flag-env-var

cli: Flag and env var equivalency
2022-11-20 07:51:38 -08:00 · 2022-11-20 07:51:38 -08:00 · 8b5e93eb2a
commit 8b5e93eb2a
parent 43947e518a c8739d8147
9 changed files with 314 additions and 82 deletions
--- a/ci/release/changelogs/next.md
+++ b/ci/release/changelogs/next.md
@ -6,7 +6,8 @@ For v0.0.99 we focused on X, Y and Z. Enjoy!

 #### Improvements 🔧

- Add table columns indices in edges between SQL Tables so that layout engines can route exactly between them
+- Equivalency between flags and environment variables. You can set either one for all
+  options (flags take precedence).

 #### Bugfixes 🔴

--- a/ci/release/template/man/d2.1
+++ b/ci/release/template/man/d2.1
@ -36,10 +36,22 @@ See more docs, the source code and license at
 .It Fl w , -watch Ar false
 Watch for changes to input and live reload. Use
 .Ev $PORT and Ev $HOST to specify the listening address.
-.Ev $D2_PORT and $D2_HOST are also accepted and take priority. Default is localhost:0
+.It Fl h , -host Ar localhost
+Host listening address when used with
+.Ar watch
+.Ns .
+.It Fl p , -port Ar 0
+Port listening address when used with
+.Ar watch
+.Ns .
 .It Fl t , -theme Ar 0
 Set the diagram theme to the passed integer. For a list of available options, see
 .Lk https://oss.terrastruct.com/d2
+.Ns .
+.It Fl l , -layout Ar dagre
+Set the diagram layout engine to the passed string. For a list of available options, run
+.Ar layout
+.Ns .
 .It Fl b , -bundle Ar true
 Bundle all assets and layers into the output svg.
 .It Fl d , -debug
--- a/cmd/d2/help.go
+++ b/cmd/d2/help.go
@ -28,13 +28,13 @@ Subcommands:
  %[1]s layout [layout name] - Display long help for a particular layout engine

 See more docs and the source code at https://oss.terrastruct.com/d2
-`, ms.Name, ms.FlagHelp())
+`, ms.Name, ms.Opts.Defaults())
 }

 func layoutHelp(ctx context.Context, ms *xmain.State) error {
-	if len(ms.FlagSet.Args()) == 1 {
+	if len(ms.Opts.Flags.Args()) == 1 {
 		return shortLayoutHelp(ctx, ms)
-	} else if len(ms.FlagSet.Args()) == 2 {
+	} else if len(ms.Opts.Flags.Args()) == 2 {
 		return longLayoutHelp(ctx, ms)
 	} else {
 		return pluginSubcommand(ctx, ms)
@ -61,7 +61,7 @@ func shortLayoutHelp(ctx context.Context, ms *xmain.State) error {
 %s

 Usage:
-  To use a particular layout engine, set the environment variable D2_LAYOUT=[layout name].
+  To use a particular layout engine, set the environment variable D2_LAYOUT=[name] or flag --layout=[name].

 Example:
  D2_LAYOUT=dagre d2 in.d2 out.svg
@ -75,7 +75,7 @@ See more docs at https://oss.terrastruct.com/d2
 }

 func longLayoutHelp(ctx context.Context, ms *xmain.State) error {
-	layout := ms.FlagSet.Arg(1)
+	layout := ms.Opts.Flags.Arg(1)
 	plugin, path, err := d2plugin.FindPlugin(ctx, layout)
 	if errors.Is(err, exec.ErrNotFound) {
 		return layoutNotFound(ctx, layout)
@ -119,13 +119,13 @@ For more information on setup, please visit https://github.com/terrastruct/d2.`,
 }

 func pluginSubcommand(ctx context.Context, ms *xmain.State) error {
-	layout := ms.FlagSet.Arg(1)
+	layout := ms.Opts.Flags.Arg(1)
 	plugin, _, err := d2plugin.FindPlugin(ctx, layout)
 	if errors.Is(err, exec.ErrNotFound) {
 		return layoutNotFound(ctx, layout)
 	}

-	ms.Args = ms.FlagSet.Args()[2:]
+	ms.Opts.Args = ms.Opts.Flags.Args()[2:]
 	return d2plugin.Serve(plugin)(ctx, ms)
 }

--- a/cmd/d2/main.go
+++ b/cmd/d2/main.go
@ -6,7 +6,6 @@ import (
 	"fmt"
 	"os/exec"
 	"path/filepath"
-	"strconv"
 	"strings"
 	"time"

@ -32,19 +31,38 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	// :(
 	ctx = xmain.DiscardSlog(ctx)

-	watchFlag := ms.FlagSet.BoolP("watch", "w", false, "watch for changes to input and live reload. Use $PORT and $HOST to specify the listening address.\n$D2_PORT and $D2_HOST are also accepted and take priority. Default is localhost:0")
-	themeFlag := ms.FlagSet.Int64P("theme", "t", 0, "set the diagram theme. For a list of available options, see https://oss.terrastruct.com/d2")
-	bundleFlag := ms.FlagSet.BoolP("bundle", "b", true, "bundle all assets and layers into the output svg")
-	versionFlag := ms.FlagSet.BoolP("version", "v", false, "get the version")
-	debugFlag := ms.FlagSet.BoolP("debug", "d", false, "print debug logs")
-	err = ms.FlagSet.Parse(ms.Args)
+	// These should be kept up-to-date with the d2 man page
+	watchFlag, err := ms.Opts.Bool("D2_WATCH", "watch", "w", false, "watch for changes to input and live reload. Use $HOST and $PORT to specify the listening address.\n(default localhost:0, which is will open on a randomly available local port).")
+	if err != nil {
+		return xmain.UsageErrorf(err.Error())
+	}
+	hostFlag := ms.Opts.String("HOST", "host", "h", "localhost", "host listening address when used with watch")
+	portFlag := ms.Opts.String("PORT", "port", "p", "0", "port listening address when used with watch")
+	bundleFlag, err := ms.Opts.Bool("D2_BUNDLE", "bundle", "b", true, "bundle all assets and layers into the output svg.")
+	if err != nil {
+		return xmain.UsageErrorf(err.Error())
+	}
+	debugFlag, err := ms.Opts.Bool("DEBUG", "debug", "d", false, "print debug logs.")
+	if err != nil {
+		return xmain.UsageErrorf(err.Error())
+	}
+	layoutFlag := ms.Opts.String("D2_LAYOUT", "layout", "l", "dagre", `the layout engine used.`)
+	themeFlag, err := ms.Opts.Int64("D2_THEME", "theme", "t", 0, "the diagram theme ID. For a list of available options, see https://oss.terrastruct.com/d2")
+	if err != nil {
+		return xmain.UsageErrorf(err.Error())
+	}
+	versionFlag, err := ms.Opts.Bool("", "version", "v", false, "get the version")
+	if err != nil {
+		return xmain.UsageErrorf(err.Error())
+	}

+	err = ms.Opts.Flags.Parse(ms.Opts.Args)
 	if !errors.Is(err, pflag.ErrHelp) && err != nil {
 		return xmain.UsageErrorf("failed to parse flags: %v", err)
 	}

-	if len(ms.FlagSet.Args()) > 0 {
-		switch ms.FlagSet.Arg(0) {
+	if len(ms.Opts.Flags.Args()) > 0 {
+		switch ms.Opts.Flags.Arg(0) {
 		case "layout":
 			return layoutHelp(ctx, ms)
 		}
@ -62,25 +80,26 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	var inputPath string
 	var outputPath string

-	if len(ms.FlagSet.Args()) == 0 {
+	if len(ms.Opts.Flags.Args()) == 0 {
 		if versionFlag != nil && *versionFlag {
 			fmt.Println(version.Version)
 			return nil
 		}
 		help(ms)
 		return nil
-	} else if len(ms.FlagSet.Args()) >= 3 {
+	} else if len(ms.Opts.Flags.Args()) >= 3 {
 		return xmain.UsageErrorf("too many arguments passed")
 	}
-	if len(ms.FlagSet.Args()) >= 1 {
-		if ms.FlagSet.Arg(0) == "version" {
+
+	if len(ms.Opts.Flags.Args()) >= 1 {
+		if ms.Opts.Flags.Arg(0) == "version" {
 			fmt.Println(version.Version)
 			return nil
 		}
-		inputPath = ms.FlagSet.Arg(0)
+		inputPath = ms.Opts.Flags.Arg(0)
 	}
-	if len(ms.FlagSet.Args()) >= 2 {
-		outputPath = ms.FlagSet.Arg(1)
+	if len(ms.Opts.Flags.Args()) >= 2 {
+		outputPath = ms.Opts.Flags.Arg(1)
 	} else {
 		if inputPath == "-" {
 			outputPath = "-"
@ -93,16 +112,11 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	if match == (d2themes.Theme{}) {
 		return xmain.UsageErrorf("-t[heme] could not be found. The available options are:\n%s\nYou provided: %d", d2themescatalog.CLIString(), *themeFlag)
 	}
-	ms.Env.Setenv("D2_THEME", fmt.Sprintf("%d", *themeFlag))
+	ms.Log.Debug.Printf("using theme %s (ID: %d)", match.Name, *themeFlag)

-	envD2Layout := ms.Env.Getenv("D2_LAYOUT")
-	if envD2Layout == "" {
-		envD2Layout = "dagre"
-	}
-
-	plugin, path, err := d2plugin.FindPlugin(ctx, envD2Layout)
+	plugin, path, err := d2plugin.FindPlugin(ctx, *layoutFlag)
 	if errors.Is(err, exec.ErrNotFound) {
-		return layoutNotFound(ctx, envD2Layout)
+		return layoutNotFound(ctx, *layoutFlag)
 	} else if err != nil {
 		return err
 	}
@ -111,14 +125,22 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	if path != "" {
 		pluginLocation = fmt.Sprintf("executable plugin at %s", humanPath(path))
 	}
-	ms.Log.Debug.Printf("using layout plugin %s (%s)", envD2Layout, pluginLocation)
+	ms.Log.Debug.Printf("using layout plugin %s (%s)", *layoutFlag, pluginLocation)

 	if *watchFlag {
 		if inputPath == "-" {
 			return xmain.UsageErrorf("-w[atch] cannot be combined with reading input from stdin")
 		}
 		ms.Env.Setenv("LOG_TIMESTAMPS", "1")
-		w, err := newWatcher(ctx, ms, plugin, inputPath, outputPath)
+
+		w, err := newWatcher(ctx, ms, watcherOpts{
+			layoutPlugin: plugin,
+			themeID:      *themeFlag,
+			host:         *hostFlag,
+			port:         *portFlag,
+			inputPath:    inputPath,
+			outputPath:   outputPath,
+		})
 		if err != nil {
 			return err
 		}
@ -132,7 +154,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 		_ = 343
 	}

-	_, err = compile(ctx, ms, plugin, inputPath, outputPath)
+	_, err = compile(ctx, ms, plugin, *themeFlag, inputPath, outputPath)
 	if err != nil {
 		return err
 	}
@ -140,7 +162,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
 	return nil
 }

-func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, inputPath, outputPath string) ([]byte, error) {
+func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, themeID int64, inputPath, outputPath string) ([]byte, error) {
 	input, err := ms.ReadPath(inputPath)
 	if err != nil {
 		return nil, err
@ -151,7 +173,6 @@ func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, input
 		return nil, err
 	}

-	themeID, _ := strconv.ParseInt(ms.Env.Getenv("D2_THEME"), 10, 64)
 	d, err := d2.Compile(ctx, string(input), &d2.CompileOptions{
 		Layout:  plugin.Layout,
 		Ruler:   ruler,
--- a/cmd/d2/watch.go
+++ b/cmd/d2/watch.go
@ -34,6 +34,15 @@ var devMode = false
 //go:embed static
 var staticFS embed.FS

+type watcherOpts struct {
+	layoutPlugin d2plugin.Plugin
+	themeID      int64
+	host         string
+	port         string
+	inputPath    string
+	outputPath   string
+}
+
 type watcher struct {
 	ctx     context.Context
 	cancel  context.CancelFunc
@ -41,9 +50,7 @@ type watcher struct {
 	devMode bool

 	ms *xmain.State
-	layoutPlugin d2plugin.Plugin
-	inputPath    string
-	outputPath   string
+	watcherOpts

 	compileCh chan struct{}

@ -68,7 +75,7 @@ type compileResult struct {
 	SVG string `json:"svg"`
 }

-func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plugin, inputPath, outputPath string) (*watcher, error) {
+func newWatcher(ctx context.Context, ms *xmain.State, opts watcherOpts) (*watcher, error) {
 	ctx, cancel := context.WithCancel(ctx)

 	w := &watcher{
@ -77,9 +84,7 @@ func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plug
 		devMode: devMode,

 		ms:          ms,
-		layoutPlugin: layoutPlugin,
-		inputPath:    inputPath,
-		outputPath:   outputPath,
+		watcherOpts: opts,

 		compileCh: make(chan struct{}, 1),
 		wsclients: make(map[*wsclient]struct{}),
@ -325,7 +330,7 @@ func (w *watcher) compileLoop(ctx context.Context) error {
 			recompiledPrefix = "re"
 		}

-		b, err := compile(ctx, w.ms, w.layoutPlugin, w.inputPath, w.outputPath)
+		b, err := compile(ctx, w.ms, w.layoutPlugin, w.themeID, w.inputPath, w.outputPath)
 		if err != nil {
 			err = fmt.Errorf("failed to %scompile: %w", recompiledPrefix, err)
 			w.ms.Log.Error.Print(err)
@ -351,18 +356,7 @@ func (w *watcher) compileLoop(ctx context.Context) error {
 }

 func (w *watcher) listen() error {
-	host := "localhost"
-	port := "0"
-	hostEnv := w.ms.Env.Getenv("HOST")
-	if hostEnv != "" {
-		host = hostEnv
-	}
-	portEnv := w.ms.Env.Getenv("PORT")
-	if portEnv != "" {
-		port = portEnv
-	}
-
-	l, err := net.Listen("tcp", net.JoinHostPort(host, port))
+	l, err := net.Listen("tcp", net.JoinHostPort(w.host, w.port))
 	if err != nil {
 		return err
 	}
--- a/d2plugin/serve.go
+++ b/d2plugin/serve.go
@ -19,12 +19,12 @@ import (
 // Also see execPlugin in exec.go for the d2 binary plugin protocol.
 func Serve(p Plugin) func(context.Context, *xmain.State) error {
 	return func(ctx context.Context, ms *xmain.State) (err error) {
-		if len(ms.Args) < 1 {
+		if len(ms.Opts.Flags.Args()) < 1 {
 			return errors.New("expected first argument to plugin binary to be function name")
 		}
-		reqFunc := ms.Args[0]
+		reqFunc := ms.Opts.Flags.Arg(0)

-		switch ms.Args[0] {
+		switch ms.Opts.Flags.Arg(0) {
 		case "info":
 			return info(ctx, p, ms)
 		case "layout":
--- a/lib/xmain/flag_helpers.go
+++ b/lib/xmain/flag_helpers.go
@ -0,0 +1,45 @@
+// flag_helpers.go are private functions from pflag/flag.go
+package xmain
+
+import "strings"
+
+func wrap(i, w int, s string) string {
+	if w == 0 {
+		return strings.Replace(s, "\n", "\n"+strings.Repeat(" ", i), -1)
+	}
+	wrap := w - i
+	var r, l string
+	if wrap < 24 {
+		i = 16
+		wrap = w - i
+		r += "\n" + strings.Repeat(" ", i)
+	}
+	if wrap < 24 {
+		return strings.Replace(s, "\n", r, -1)
+	}
+	slop := 5
+	wrap = wrap - slop
+	l, s = wrapN(wrap, slop, s)
+	r = r + strings.Replace(l, "\n", "\n"+strings.Repeat(" ", i), -1)
+	for s != "" {
+		var t string
+		t, s = wrapN(wrap, slop, s)
+		r = r + "\n" + strings.Repeat(" ", i) + strings.Replace(t, "\n", "\n"+strings.Repeat(" ", i), -1)
+	}
+	return r
+}
+
+func wrapN(i, slop int, s string) (string, string) {
+	if i+slop > len(s) {
+		return s, ""
+	}
+	w := strings.LastIndexAny(s[:i], " \t\n")
+	if w <= 0 {
+		return s, ""
+	}
+	nlPos := strings.LastIndex(s[:i], "\n")
+	if nlPos > 0 && nlPos < w {
+		return s[:nlPos], s[nlPos+1:]
+	}
+	return s[:w], s[w+1:]
+}
--- a/lib/xmain/opts.go
+++ b/lib/xmain/opts.go
@ -0,0 +1,173 @@
+package xmain
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"strconv"
+	"strings"
+
+	"github.com/spf13/pflag"
+	"oss.terrastruct.com/cmdlog"
+	"oss.terrastruct.com/xos"
+)
+
+type Opts struct {
+	Args  []string
+	Flags *pflag.FlagSet
+	env   *xos.Env
+	log   *cmdlog.Logger
+
+	flagEnv map[string]string
+}
+
+func NewOpts(env *xos.Env, log *cmdlog.Logger, args []string) *Opts {
+	flags := pflag.NewFlagSet("", pflag.ContinueOnError)
+	flags.SortFlags = false
+	flags.Usage = func() {}
+	flags.SetOutput(io.Discard)
+	return &Opts{
+		Args:    args,
+		Flags:   flags,
+		env:     env,
+		log:     log,
+		flagEnv: make(map[string]string),
+	}
+}
+
+// Mostly copy pasted pasted from pflag.FlagUsagesWrapped
+// with modifications for env var
+func (o *Opts) Defaults() string {
+	buf := new(bytes.Buffer)
+
+	var lines []string
+
+	maxlen := 0
+	maxEnvLen := 0
+	o.Flags.VisitAll(func(flag *pflag.Flag) {
+		if flag.Hidden {
+			return
+		}
+
+		line := ""
+		if flag.Shorthand != "" && flag.ShorthandDeprecated == "" {
+			line = fmt.Sprintf("  -%s, --%s", flag.Shorthand, flag.Name)
+		} else {
+			line = fmt.Sprintf("      --%s", flag.Name)
+		}
+
+		varname, usage := pflag.UnquoteUsage(flag)
+		if varname != "" {
+			line += " " + varname
+		}
+		if flag.NoOptDefVal != "" {
+			switch flag.Value.Type() {
+			case "string":
+				line += fmt.Sprintf("[=\"%s\"]", flag.NoOptDefVal)
+			case "bool":
+				if flag.NoOptDefVal != "true" {
+					line += fmt.Sprintf("[=%s]", flag.NoOptDefVal)
+				}
+			case "count":
+				if flag.NoOptDefVal != "+1" {
+					line += fmt.Sprintf("[=%s]", flag.NoOptDefVal)
+				}
+			default:
+				line += fmt.Sprintf("[=%s]", flag.NoOptDefVal)
+			}
+		}
+
+		line += "\x00"
+
+		if len(line) > maxlen {
+			maxlen = len(line)
+		}
+
+		if e, ok := o.flagEnv[flag.Name]; ok {
+			line += fmt.Sprintf("$%s", e)
+		}
+
+		line += "\x01"
+
+		if len(line) > maxEnvLen {
+			maxEnvLen = len(line)
+		}
+
+		line += usage
+		if flag.Value.Type() == "string" {
+			line += fmt.Sprintf(" (default %q)", flag.DefValue)
+		} else {
+			line += fmt.Sprintf(" (default %s)", flag.DefValue)
+		}
+		if len(flag.Deprecated) != 0 {
+			line += fmt.Sprintf(" (DEPRECATED: %s)", flag.Deprecated)
+		}
+
+		lines = append(lines, line)
+	})
+
+	for _, line := range lines {
+		sidx1 := strings.Index(line, "\x00")
+		sidx2 := strings.Index(line, "\x01")
+		spacing1 := strings.Repeat(" ", maxlen-sidx1)
+		spacing2 := strings.Repeat(" ", (maxEnvLen-maxlen)-sidx2+sidx1)
+		fmt.Fprintln(buf, line[:sidx1], spacing1, line[sidx1+1:sidx2], spacing2, wrap(maxEnvLen+3, 0, line[sidx2+1:]))
+	}
+
+	return buf.String()
+}
+
+func (o *Opts) getEnv(flag, k string) string {
+	if k != "" {
+		o.flagEnv[flag] = k
+		return o.env.Getenv(k)
+	}
+	return ""
+}
+
+func (o *Opts) Int64(envKey, flag, shortFlag string, defaultVal int64, usage string) (*int64, error) {
+	if env := o.getEnv(flag, envKey); env != "" {
+		envVal, err := strconv.ParseInt(env, 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf(`invalid environment variable %s. Expected int64. Found "%v".`, envKey, envVal)
+		}
+		defaultVal = envVal
+	}
+
+	return o.Flags.Int64P(flag, shortFlag, defaultVal, usage), nil
+}
+
+func (o *Opts) String(envKey, flag, shortFlag string, defaultVal, usage string) *string {
+	if env := o.getEnv(flag, envKey); env != "" {
+		defaultVal = env
+	}
+
+	return o.Flags.StringP(flag, shortFlag, defaultVal, usage)
+}
+
+func (o *Opts) Bool(envKey, flag, shortFlag string, defaultVal bool, usage string) (*bool, error) {
+	if env := o.getEnv(flag, envKey); env != "" {
+		if !boolyEnv(env) {
+			return nil, fmt.Errorf(`invalid environment variable %s. Expected bool. Found "%s".`, envKey, env)
+		}
+		if truthyEnv(env) {
+			defaultVal = true
+		} else {
+			defaultVal = false
+		}
+	}
+
+	return o.Flags.BoolP(flag, shortFlag, defaultVal, usage), nil
+}
+
+func boolyEnv(s string) bool {
+	return falseyEnv(s) || truthyEnv(s)
+}
+
+func falseyEnv(s string) bool {
+	return s == "0" || s == "false"
+}
+
+func truthyEnv(s string) bool {
+	return s == "1" || s == "true"
+}
--- a/lib/xmain/xmain.go
+++ b/lib/xmain/xmain.go
@ -9,13 +9,11 @@ import (
 	"io"
 	"os"
 	"os/signal"
-	"strings"
 	"syscall"
 	"time"

 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
-	"github.com/spf13/pflag"

 	"oss.terrastruct.com/xos"

@ -42,13 +40,9 @@ func Main(run RunFunc) {
 		Stderr: os.Stderr,

 		Env: xos.NewEnv(os.Environ()),
-		FlagSet: pflag.NewFlagSet("", pflag.ContinueOnError),
-		Args:    args,
 	}
 	ms.Log = cmdlog.Log(ms.Env, os.Stderr)
-	ms.FlagSet.SortFlags = false
-	ms.FlagSet.Usage = func() {}
-	ms.FlagSet.SetOutput(io.Discard)
+	ms.Opts = NewOpts(ms.Env, ms.Log, args)

 	sigs := make(chan os.Signal, 1)
 	signal.Notify(sigs, os.Interrupt, syscall.SIGTERM)
@ -90,8 +84,7 @@ type State struct {

 	Log  *cmdlog.Logger
 	Env  *xos.Env
-	Args    []string
-	FlagSet *pflag.FlagSet
+	Opts *Opts
 }

 func (ms *State) Main(ctx context.Context, sigs <-chan os.Signal, run func(context.Context, *State) error) error {
@ -129,13 +122,6 @@ func (ms *State) Main(ctx context.Context, sigs <-chan os.Signal, run func(conte
 	}
 }

-func (ms *State) FlagHelp() string {
-	b := &strings.Builder{}
-	ms.FlagSet.SetOutput(b)
-	ms.FlagSet.PrintDefaults()
-	return b.String()
-}
-
 type ExitError struct {
 	Code    int    `json:"code"`
 	Message string `json:"message"`