Teach pflag to handle unambiguous abbreviated long-args

abbrev.go

@@ -0,0 +1,57 @@
+// abbrev.go - generate abbreviations from a wordlist
+//
+// (c) 2014 Sudhi Herle <sw-at-herle.net>
+//
+// Placed in the Public Domain
+// This software does not come with any express or implied
+// warranty; it is provided "as is". No claim  is made to its
+// suitability for any purpose.
+
+package pflag
+
+// Given a wordlist in 'words', generate unique abbreviations of it
+// and return as a map[abbrev]word.
+// e.g.,
+//
+//	given a wordlist ["hello", "help", "sync"],
+//	Abbrev() returns:
+//	  {
+//	     "hello": "hello",
+//	     "hell":  "hell"
+//	     "help":  "help",
+//	     "sync":  "sync",
+//	     "syn":   "sync",
+//	     "sy":    "sync",
+//	     "s":     "sync"
+//	  }
+func abbrev(words []string) map[string]string {
+	seen := make(map[string]int)
+	table := make(map[string]string)
+
+	for _, w := range words {
+		for n := len(w) - 1; n > 0; n -= 1 {
+			ab := w[:n]
+			seen[ab] += 1
+
+			switch seen[ab] {
+			case 1:
+				table[ab] = w
+			case 2:
+				delete(table, ab)
+			default:
+				goto next
+			}
+		}
+	next:
+	}
+
+	// non abbreviations always get entered
+	// This has to be done _after_ the loop above; because
+	// if there are words that are prefixes of other words in
+	// the argument list, we need to ensure we capture them
+	// intact.
+	for _, w := range words {
+		table[w] = w
+	}
+	return table
+}

abbrev_test.go

@@ -0,0 +1,60 @@
+package pflag
+
+import (
+	"fmt"
+	"runtime"
+	"testing"
+)
+
+func TestSimple(t *testing.T) {
+	assert := newAsserter(t)
+
+	words := []string{"hello", "help", "sync", "uint", "uint16", "uint64"}
+	ret := map[string]string{
+		"hello":  "hello",
+		"hell":   "hello",
+		"help":   "help",
+		"sync":   "sync",
+		"syn":    "sync",
+		"sy":     "sync",
+		"s":      "sync",
+		"uint":   "uint",
+		"uint16": "uint16",
+		"uint1":  "uint16",
+		"uint64": "uint64",
+		"uint6":  "uint64",
+	}
+
+	ab := abbrev(words)
+
+	for k, v := range ab {
+		x, ok := ret[k]
+		assert(ok, "unexpected abbrev %s", k)
+		assert(x == v, "abbrev %s: exp %s, saw %s", k, x, v)
+	}
+
+	for k, v := range ret {
+		x, ok := ab[k]
+		assert(ok, "unknown abbrev %s", k)
+		assert(x == v, "abbrev %s: exp %s, saw %s", k, x, v)
+	}
+
+}
+
+// make an assert() function for use in environment 't' and return it
+func newAsserter(t *testing.T) func(cond bool, msg string, args ...interface{}) {
+	return func(cond bool, msg string, args ...interface{}) {
+		if cond {
+			return
+		}
+
+		_, file, line, ok := runtime.Caller(1)
+		if !ok {
+			file = "???"
+			line = 0
+		}
+
+		s := fmt.Sprintf(msg, args...)
+		t.Fatalf("%s: %d: Assertion failed: %s\n", file, line, s)
+	}
+}

example_test.go

@@ -7,7 +7,7 @@ package pflag_test
 import (
 	"fmt"
 
-	"github.com/spf13/pflag"
+	"github.com/opencoff/pflag"
 )
 
 func ExampleShorthandLookup() {

flag.go

@@ -16,7 +16,7 @@ pflag is a drop-in replacement of Go's native flag package. If you import
 pflag under the name "flag" then all code should continue to function
 with no changes.
 
-	import flag "github.com/spf13/pflag"
+	import flag "github.com/opencoff/pflag"
 
 There is one exception to this: if you directly instantiate the Flag struct
 there is one more field "Shorthand" that you will need to set.
@@ -27,23 +27,32 @@ unaffected.
 Define flags using flag.String(), Bool(), Int(), etc.
 
 This declares an integer flag, -flagname, stored in the pointer ip, with type *int.
+
 	var ip = flag.Int("flagname", 1234, "help message for flagname")
+
 If you like, you can bind the flag to a variable using the Var() functions.
+
 	var flagvar int
 	func init() {
 		flag.IntVar(&flagvar, "flagname", 1234, "help message for flagname")
 	}
+
 Or you can create custom flags that satisfy the Value interface (with
 pointer receivers) and couple them to flag parsing by
+
 	flag.Var(&flagVal, "name", "help message for flagname")
+
 For such flags, the default value is just the initial value of the variable.
 
 After all flags are defined, call
+
 	flag.Parse()
+
 to parse the command line into the defined flags.
 
 Flags may then be used directly. If you're using the flags themselves,
 they are all pointers; if you bind to variables, they're values.
+
 	fmt.Println("ip has value ", *ip)
 	fmt.Println("flagvar has value ", flagvar)
 
@@ -54,22 +63,26 @@ The arguments are indexed from 0 through flag.NArg()-1.
 The pflag package also defines some new functions that are not in flag,
 that give one-letter shorthands for flags. You can use these by appending
 'P' to the name of any function that defines a flag.
+
 	var ip = flag.IntP("flagname", "f", 1234, "help message")
 	var flagvar bool
 	func init() {
 		flag.BoolVarP(&flagvar, "boolname", "b", true, "help message")
 	}
 	flag.VarP(&flagval, "varname", "v", "help message")
+
 Shorthand letters can be used with single dashes on the command line.
 Boolean shorthand flags can be combined with other shorthand flags.
 
 Command line flag syntax:
+
 	--flag    // boolean flags only
 	--flag=x
 
 Unlike the flag package, a single dash before an option means something
 different than a double dash. Single dashes signify a series of shorthand
 letters for flags. All but the last shorthand letter must be boolean flags.
+
 	// boolean flags
 	-f
 	-abc
@@ -84,6 +97,26 @@ Flag parsing stops after the terminator "--". Unlike the flag package,
 flags can be interspersed with arguments anywhere on the command line
 before this terminator.
 
+Long form flags can also be abbreviated - so long as it is a unique
+abbreviation. eg given this:
+
+	var ip = flag.IntP("flagname", "f", 1234, "help message")
+	var op = flag.IntP("fluid-level", "F", 99, "fluid message")
+
+The following abbreviations will all resolve to "flagname":
+
+	--flagname=33
+	--flag=33
+	--fla=33
+	-f 33
+
+And the following abbreviations will all resolve to "fluid-level":
+
+	--fluid-level=20
+	--fluid=20
+	--flu=20
+	-F 20
+
 Integer flags accept 1234, 0664, 0x1234 and may be negative.
 Boolean flags (in their long form) accept 1, 0, t, f, true, false,
 TRUE, FALSE, True, False.
@@ -165,6 +198,11 @@ type FlagSet struct {
 	normalizeNameFunc func(f *FlagSet, name string) NormalizedName
 
 	addedGoFlagSets []*goflag.FlagSet
+
+	// map to hold unambiguous, abbreviated long-name
+	// key = abbreviation;
+	// value = full arg name
+	abbrev map[string]string
 }
 
 // A Flag represents the state of a flag.
@@ -934,9 +972,9 @@ func (f *FlagSet) usage() {
 	}
 }
 
-//--unknown (args will be empty)
-//--unknown --next-flag ... (args will be --next-flag ...)
-//--unknown arg ... (args will be arg ...)
+// --unknown (args will be empty)
+// --unknown --next-flag ... (args will be --next-flag ...)
+// --unknown arg ... (args will be arg ...)
 func stripUnknownFlagValue(args []string) []string {
 	if len(args) == 0 {
 		//--unknown
@@ -965,7 +1003,15 @@ func (f *FlagSet) parseLongArg(s string, args []string, fn parseFunc) (a []strin
 	}
 
 	split := strings.SplitN(name, "=", 2)
-	name = split[0]
+	ab := split[0]
+	name, ok := f.abbrev[ab]
+	if !ok {
+		if ab == "help" {
+			f.usage()
+			return a, ErrHelp
+		}
+		name = ab
+	}
 	flag, exists := f.formal[f.normalizeFlagName(name)]
 
 	if !exists {
@@ -1123,6 +1169,16 @@ func (f *FlagSet) parseArgs(args []string, fn parseFunc) (err error) {
 	return
 }
 
+// setup abbreviations for long args
+func (f *FlagSet) setupAbbrev() {
+	// create unique shortcuts for the long args
+	words := make([]string, 0, len(f.formal))
+	for k := range f.formal {
+		words = append(words, string(k))
+	}
+	f.abbrev = abbrev(words)
+}
+
 // Parse parses flag definitions from the argument list, which should not
 // include the command name.  Must be called after all flags in the FlagSet
 // are defined and before flags are accessed by the program.
@@ -1139,6 +1195,7 @@ func (f *FlagSet) Parse(arguments []string) error {
 		return nil
 	}
 
+	f.setupAbbrev()
 	f.args = make([]string, 0, len(arguments))
 
 	set := func(flag *Flag, value string) error {
@@ -1168,6 +1225,7 @@ type parseFunc func(flag *Flag, value string) error
 // accessed by the program. The return value will be ErrHelp if -help was set
 // but not defined.
 func (f *FlagSet) ParseAll(arguments []string, fn func(flag *Flag, value string) error) error {
+	f.setupAbbrev()
 	f.parsed = true
 	f.args = make([]string, 0, len(arguments))
 

flag_test.go

@@ -1134,7 +1134,6 @@ func TestMultipleNormalizeFlagNameInvocations(t *testing.T) {
 	}
 }
 
-//
 func TestHiddenFlagInUsage(t *testing.T) {
 	f := NewFlagSet("bob", ContinueOnError)
 	f.Bool("secretFlag", true, "shhh")
@@ -1149,7 +1148,6 @@ func TestHiddenFlagInUsage(t *testing.T) {
 	}
 }
 
-//
 func TestHiddenFlagUsage(t *testing.T) {
 	f := NewFlagSet("bob", ContinueOnError)
 	f.Bool("secretFlag", true, "shhh")
@@ -1238,16 +1236,16 @@ func TestPrintDefaults(t *testing.T) {
 	fs.PrintDefaults()
 	got := buf.String()
 	if got != defaultOutput {
-		fmt.Println("\n" + got)
-		fmt.Println("\n" + defaultOutput)
+		fmt.Print("\n" + got)
+		fmt.Print("\n" + defaultOutput)
 		t.Errorf("got %q want %q\n", got, defaultOutput)
 	}
 }
 
 func TestVisitAllFlagOrder(t *testing.T) {
 	fs := NewFlagSet("TestVisitAllFlagOrder", ContinueOnError)
 	fs.SortFlags = false
-	// https://github.com/spf13/pflag/issues/120
+	// https://github.com/opencoff/pflag/issues/120
 	fs.SetNormalizeFunc(func(f *FlagSet, name string) NormalizedName {
 		return NormalizedName(name)
 	})
@@ -1283,3 +1281,32 @@ func TestVisitFlagOrder(t *testing.T) {
 		i++
 	})
 }
+
+func TestAbbrevFlags(t *testing.T) {
+	f := NewFlagSet("abbrev", ContinueOnError)
+
+	b0 := f.Bool("with-something", false, "bool with something")
+	b1 := f.Bool("with-otherthing", false, "bool with otherthing")
+	b2 := f.Bool("zero-this", false, "zero this thing")
+
+	args := []string{
+		"--with-some",
+		"--with-oth",
+		"--ze",
+	}
+	if err := f.Parse(args); err != nil {
+		t.Fatal(err)
+	}
+	if !f.Parsed() {
+		t.Error("f.Parse() = false after Parse")
+	}
+	if *b0 != true {
+		t.Error("with-something flag should be true, is ", *b0)
+	}
+	if *b1 != true {
+		t.Error("with-otherthing flag should be true, is ", *b1)
+	}
+	if *b2 != true {
+		t.Error("zero-this flag should be true, is ", *b2)
+	}
+}

go.mod

@@ -1,3 +1,3 @@
-module github.com/spf13/pflag
+module github.com/opencoff/pflag
 
-go 1.12
+go 1.20