From d480572d75eb82a6052e1eb7d16268983820608f Mon Sep 17 00:00:00 2001
From: Alec Thomas <alec@swapoff.org>
Date: Wed, 10 Jun 2020 15:40:22 +1000
Subject: [PATCH] Make counter flags more flexible.

They now support three forms:

    -s
    --long
    --long=N

The last of which explicitly sets the counter value.

Fixes #87.
---
 README.md      |  2 +-
 mapper.go      | 22 ++++++++++++++++++++++
 mapper_test.go | 10 +++++++++-
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 508a1df..eaad200 100644
--- a/README.md
+++ b/README.md
@@ -375,7 +375,7 @@ function `NamedMapper(name, mapper)`.
 | `path`            | A path. ~ expansion is applied.          
 | `existingfile`    | An existing file. ~ expansion is applied. `-` is accepted for stdin.
 | `existingdir`     | An existing directory. ~ expansion is applied.   
-| `counter`         | Increment a numeric field. Useful for `-vvv`    
+| `counter`         | Increment a numeric field. Useful for `-vvv`. Can accept `-s`, `--long` or `--long=N`.
 
 
 Slices and maps treat type tags specially. For slices, the `type:""` tag
diff --git a/mapper.go b/mapper.go
index 8a1944d..b0f99b1 100644
--- a/mapper.go
+++ b/mapper.go
@@ -597,6 +597,28 @@ func existingDirMapper(r *Registry) MapperFunc {
 
 func counterMapper() MapperFunc {
 	return func(ctx *DecodeContext, target reflect.Value) error {
+		if ctx.Scan.Peek().Type == FlagValueToken {
+			t, err := ctx.Scan.PopValue("counter")
+			if err != nil {
+				return err
+			}
+			switch v := t.Value.(type) {
+			case string:
+				n, err := strconv.ParseInt(v, 10, 64)
+				if err != nil {
+					return errors.Errorf("expected a counter but got %q (%T)", t, t.Value)
+				}
+				target.SetInt(n)
+
+			case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
+				target.Set(reflect.ValueOf(v))
+
+			default:
+				return errors.Errorf("expected a counter but got %q (%T)", t, t.Value)
+			}
+			return nil
+		}
+
 		switch target.Kind() {
 		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
 			target.SetInt(target.Int() + 1)
diff --git a/mapper_test.go b/mapper_test.go
index b6d7737..c0ba855 100644
--- a/mapper_test.go
+++ b/mapper_test.go
@@ -262,7 +262,15 @@ func TestCounter(t *testing.T) {
 	}
 	p := mustNew(t, &cli)
 
-	_, err := p.Parse([]string{"-iii"})
+	_, err := p.Parse([]string{"--int", "--int", "--int"})
+	require.NoError(t, err)
+	require.Equal(t, 3, cli.Int)
+
+	_, err = p.Parse([]string{"--int=5"})
+	require.NoError(t, err)
+	require.Equal(t, 5, cli.Int)
+
+	_, err = p.Parse([]string{"-iii"})
 	require.NoError(t, err)
 	require.Equal(t, 3, cli.Int)