Support optional-args for short options

2 files changed, 58 insertions, 10 deletions

diff --git a/opts.go b/opts.go
index 3d07b0e..4fe07a5 100644
--- a/opts.go
+++ b/opts.go
@@ -62,14 +62,23 @@ type LongOpt struct {
 // of the first non-option argument in optind.  In the case of failure,
 // err will be one of [BadOptionError] or [NoArgumentError].
 func Get(args []string, optstr string) (flags []Flag, optind int, err error) {
-	argmap := make(map[rune]bool)
+	argmap := make(map[rune]ArgMode)
 
 	rs := []rune(optstr)
-	for i, r := range rs {
-		if r != ':' {
-			argmap[r] = false
-		} else if i > 0 {
-			argmap[rs[i-1]] = true
+	if rs[0] == ':' {
+		rs = rs[1:]
+	}
+	for len(rs) > 0 {
+		switch r := rs[0]; {
+		case len(rs) > 2 && rs[1] == ':' && rs[2] == ':':
+			argmap[r] = Optional
+			rs = rs[3:]
+		case len(rs) > 1 && rs[1] == ':':
+			argmap[r] = Required
+			rs = rs[2:]
+		default:
+			argmap[r] = None
+			rs = rs[1:]
 		}
 	}
 
@@ -86,14 +95,14 @@ func Get(args []string, optstr string) (flags []Flag, optind int, err error) {
 		rs := []rune(arg[1:])
 		for j, r := range rs {
 			var s string
-			argp, ok := argmap[r]
+			am, ok := argmap[r]
 
 			switch {
 			case !ok:
 				return nil, 0, BadOptionError(r)
-			case argp && j < len(rs)-1:
+			case am != None && j < len(rs)-1:
 				s = string(rs[j+1:])
-			case argp:
+			case am == Required:
 				i++
 				if i >= len(args) {
 					return nil, 0, NoArgumentError(r)
diff --git a/opts_test.go b/opts_test.go
index 2ea9334..1f71057 100644
--- a/opts_test.go
+++ b/opts_test.go
@@ -10,7 +10,7 @@ func die(t *testing.T, name string, want, got any) {
 }
 
 func assertGet(t *testing.T, args []string, fw, ow int, ew error) []Flag {
-	flags, optind, err := Get(args, "abλc:dß")
+	flags, optind, err := Get(args, "abλc:dßĦ::")
 	if err != ew {
 		die(t, "err", ew, err)
 	}
@@ -217,3 +217,42 @@ func TestXAfterDash(t *testing.T) {
 	args := []string{"foo", "-", "-x"}
 	assertGet(t, args, 0, 1, nil)
 }
+
+func TestĦWithSpaceAndArg(t *testing.T) {
+	args := []string{"foo", "-Ħ", "bar"}
+	flags := assertGet(t, args, 1, 2, nil)
+	if flags[0].Key != 'Ħ' {
+		die(t, "flags[0].Key", 'Ħ', flags[0].Key)
+	}
+	if flags[0].Value != "" {
+		die(t, "flags[0].Value", "", flags[0].Value)
+	}
+}
+
+func TestĦWithArg(t *testing.T) {
+	args := []string{"foo", "-Ħbar"}
+	flags := assertGet(t, args, 1, 2, nil)
+	if flags[0].Key != 'Ħ' {
+		die(t, "flags[0].Key", 'Ħ', flags[0].Key)
+	}
+	if flags[0].Value != "bar" {
+		die(t, "flags[0].Value", "bar", flags[0].Value)
+	}
+}
+
+func TestΛĦWithArg(t *testing.T) {
+	args := []string{"foo", "-λĦbar"}
+	flags := assertGet(t, args, 2, 2, nil)
+	if flags[0].Key != 'λ' {
+		die(t, "flags[0].Key", 'λ', flags[0].Key)
+	}
+	if flags[0].Value != "" {
+		die(t, "flags[0].Value", "", flags[0].Value)
+	}
+	if flags[1].Key != 'Ħ' {
+		die(t, "flags[1].Key", 'Ħ', flags[1].Key)
+	}
+	if flags[1].Value != "bar" {
+		die(t, "flags[1].Value", "bar", flags[1].Value)
+	}
+}