Boring multi-command dd-like argument parsing

args.go

// args.go is a CLI arg parser for ipexec and other tools I've written but not released.
// It is here mainly as a backup since I don't feel particularly compelled to make a package out of it.

// TODO: Figure out which version is newest and update this.

package main

import (
	"net/url"
	"strconv"
	"strings"
	"unicode/utf8"
)

const argSep = '='

const ArgLine = "ARGV"

type Arg struct {
	Cmd  string
	Args map[string][]string
}

// HasOnly returns true if the Arg only has arguments with the given names.
// False is returned only if the Arg has an argument with a name not present in argnames.
func (a *Arg) HasOnly(names ...string) bool {
	set := make(map[string]struct{}, len(names))
	for _, name := range names {
		set[name] = struct{}{}
	}
	for name := range a.Args {
		if _, ok := set[name]; !ok {
			return false
		}
	}
	return true
}

func (a *Arg) Has(argname string) bool {
	v := a.Args[argname]
	return len(v) > 0
}

func (a *Arg) Str(argname, def string) string {
	v := a.Args[argname]
	if len(v) > 0 {
		return v[len(v)-1]
	}
	return def
}

func (a *Arg) Int(argname string, def int) (int, error) {
	v := a.Args[argname]
	if len(v) > 0 {
		return strconv.Atoi(v[len(v)-1])
	}
	return def, nil
}

func (a *Arg) URL(argname, def string) (*url.URL, error) {
	v := a.Args[argname]
	if len(v) > 0 {
		return url.Parse(v[len(v)-1])
	}
	if def == "" {
		return nil, nil
	}
	return url.Parse(def)
}

type argType uint

const (
	argNone argType = iota
	argKVPair
	argShort
	argLong
)

type ArgMapping struct {
	typ     argType
	cmd     string
	argname string
	remapTo string
}

func RemapShort(cmd string, short rune, long string) ArgMapping {
	shortArg := string(short)
	return ArgMapping{
		typ:     argShort,
		cmd:     cmd,
		argname: shortArg,
		remapTo: long,
	}
}

func RemapKey(cmd string, name, newName string) ArgMapping {
	return ArgMapping{
		typ:     argKVPair,
		cmd:     cmd,
		argname: name,
		remapTo: newName,
	}
}

func argLineMarker(s string) bool {
	const prefix = "--"
	const markSep = '-'
	if !strings.HasPrefix(s, prefix) {
		return false
	}
	for i := len(prefix); i < len(s); i++ { // Don't bother decoding runes here
		if s[i] != markSep {
			return false
		}
	}
	return true
}

func shortArg(s string) (k, v string, ok bool) {
	if len(s) < 2 || s[0] != '-' || s[1] == '-' {
		return k, v, false
	}
	s = s[1:]
	code, sz := utf8.DecodeRuneInString(s)
	if code == utf8.RuneError {
		return k, v, false
	}
	return s[:sz], s[sz:], true
}

type argKey struct {
	typ     argType
	cmd     string
	argname string
}

type remapFunc func(typ argType, cmd, argname string) (newArgname string)

func defaultRemap(typ argType, _, argname string) string {
	switch typ {
	case argKVPair:
		return strings.ToLower(argname)
	default:
		return argname
	}
}

func remappingTableFunc(mappings []ArgMapping) remapFunc {
	if len(mappings) == 0 {
		return defaultRemap
	}
	table := make(map[argKey]string, len(mappings))
	for _, mapping := range mappings {
		k := argKey{typ: mapping.typ, cmd: mapping.cmd, argname: mapping.argname}
		table[k] = mapping.remapTo
	}
	return func(typ argType, cmd, argname string) string {
		if v, ok := table[argKey{typ: typ, cmd: cmd, argname: argname}]; ok {
			argname = v
		} else {
			argname = defaultRemap(typ, cmd, argname)
		}
		return argname
	}
}

func ParseArgs(args []string, freeformOnly bool, mappings ...ArgMapping) []Arg {
	// TODO(ncower): Add command tree parsing to bail early if arguments don't follow in the correct order.
	// TODO(ncower): Add known-argument checking to bail early if arguments aren't recognized.
	var (
		set      = make([]Arg, 0, 2)
		inv      = Arg{Args: map[string][]string{}}
		freeform bool
		freesep  string
		remap    = remappingTableFunc(mappings)
		k, v     string
		typ      argType
		ok       bool
	)
	for _, arg := range args {
		if freeform {
			if arg == freesep {
				freeform = false
				continue
			}
			inv.Args[ArgLine] = append(inv.Args[ArgLine], arg)
			continue
		} else if argLineMarker(arg) {
			freeform, freesep = true, arg
			continue
		}

		if k, v, ok = shortArg(arg); ok {
			typ = argShort
			goto consumeArg
		}

		// TODO(ncower): Support --long-args? (Personally don't like these)
		// if k, v, ok = longArg(arg); ok {
		// 	typ = argLong
		// 	goto consumeArg
		// }

		// K=V arg
		if sep := strings.IndexByte(arg, argSep); sep > -1 {
			typ = argKVPair
			k, v = arg[:sep], arg[sep+1:]
			goto consumeArg
		}

		if freeformOnly {
			typ, k, v = argNone, ArgLine, arg
			goto consumeArg
		}

		// New command
		set = append(set, inv)
		inv = Arg{
			Cmd:  strings.ToLower(arg),
			Args: map[string][]string{},
		}
		continue

	consumeArg:
		k = remap(typ, inv.Cmd, k)
		inv.Args[k] = append(inv.Args[k], v)
		continue
	}
	set = append(set, inv)
	return set
}

args_test.go

package main

import "testing"

func stringMapsEqual(l, r map[string][]string) bool {
	if len(l) != len(r) {
		return false
	}
	for k, lv := range l {
		rv := r[k]
		if !stringsEqual(lv, rv) {
			return false
		}
	}
	return true
}

func stringsEqual(l, r []string) bool {
	// Treat nil and []string{} as logically equivalent here
	if len(l) != len(r) {
		return false
	}
	for i, s := range l {
		if s != r[i] {
			return false
		}
	}
	return true
}

func mkarg(cmd string, argv ...string) Arg {
	if len(argv)%2 != 0 {
		panic("mkarg: " + cmd + ": len(argv)%2 != 0")
	}
	args := map[string][]string{}
	for i := 0; i < len(argv); i += 2 {
		k, v := argv[i], argv[i+1]
		args[k] = append(args[k], v)
	}
	return Arg{Cmd: cmd, Args: args}
}

func TestParseArgs(t *testing.T) {
	check := func(t *testing.T, freeformOnly bool, want []Arg, mappings []ArgMapping, args []string) {
		var (
			got   = ParseArgs(args, freeformOnly, mappings...)
			gotn  = len(got)
			wantn = len(want)
		)
		if gotn != wantn {
			t.Errorf("ArgN = %d; want %d", len(got), len(want))
		}

		for i, gotarg := range got {
			if i >= wantn {
				break
			}
			wantarg := want[i]
			if gotarg.Cmd != wantarg.Cmd {
				t.Errorf("Arg[%d].Cmd = %q; want %q", i, gotarg.Cmd, wantarg.Cmd)
			}
			if !stringMapsEqual(gotarg.Args, wantarg.Args) {
				t.Errorf("Arg[%d].Args = %q; want %q", i, gotarg.Args, wantarg.Args)
			}
		}
	}
	runMultiCmd := func(name string, want []Arg, mappings []ArgMapping, args ...string) {
		t.Run(name, func(t *testing.T) { check(t, false, want, mappings, args) })
	}
	runFreeform := func(name string, want []Arg, mappings []ArgMapping, args ...string) {
		t.Run(name, func(t *testing.T) { check(t, true, want, mappings, args) })
	}

	// Test cases

	runMultiCmd("Nil-MultiCmd", []Arg{mkarg("")}, nil)
	runFreeform("Nil-Free", []Arg{mkarg("")}, nil)

	runMultiCmd("Empty-MultiCmd", []Arg{mkarg("", []string{}...)}, nil, []string{}...)
	runFreeform("Empty-Free", []Arg{mkarg("", []string{}...)}, nil, []string{}...)

	runFreeform("FreeformOnly",
		[]Arg{
			mkarg("",
				ArgLine, "freeform",
				ArgLine, "baz",
				ArgLine, "beep=birb",
				"foo", "bar",
				"v", "3",
			),
		},
		nil,
		"foo=bar", "freeform", "baz", "--", "beep=birb", "--", "-v3",
	)

	runMultiCmd("FirstAlwaysDefined",
		[]Arg{
			mkarg(""),
			mkarg("help"),
		},
		nil,
		"help",
	)

	runMultiCmd("WithArguments",
		[]Arg{
			mkarg(""),
			mkarg("tcp", "bind", "127.0.0.1:9999"),
		},
		nil,
		"TCP", "BIND=127.0.0.1:9999",
	)

	runMultiCmd("InvalidShortArg",
		[]Arg{
			mkarg(""),
			mkarg("-\xf0\x9f\x98", "foo", "bar"),
		},
		nil,
		"-\xf0\x9f\x98", "foo=bar",
	)

	runMultiCmd("UTF8ShortArgs",
		[]Arg{
			mkarg("", "\U0001F914", "thinking face"),
		},
		nil,
		"-\U0001F914thinking face",
	)

	runMultiCmd("ShortArgs",
		[]Arg{
			mkarg(""),
			mkarg("tcp", "bind", "127.0.0.1:9999", "bind", ":7777", "T", "128"),
		},
		[]ArgMapping{
			RemapShort("tcp", 'b', "bind"),
			RemapKey("tcp", "listen", "bind"),
		},
		"TCP", "-b127.0.0.1:9999", "-T128", "listen=:7777",
	)

	runMultiCmd("WithoutArguments",
		[]Arg{
			mkarg(""),
			mkarg("this"),
			mkarg("has"),
			mkarg("no"),
			mkarg("arguments"),
		},
		nil,
		"this", "has", "no", "arguments",
	)

	runMultiCmd("Freeform",
		[]Arg{
			mkarg("",
				ArgLine, "freeform",
				"--foobar", "baz",
			),
			mkarg("tcp",
				"bind", "127.0.0.1:9999",
				ArgLine, "command",
				ArgLine, "--",
				ArgLine, "--argument=foo",
				ArgLine, "/dev/urandom",
			),
		},
		nil,
		"--", "freeform", "--", "--foobar=baz",
		"tcp", "BIND=127.0.0.1:9999", "------", "command", "--", "--argument=foo", "/dev/urandom", "------",
	)
}