From 2ee16310d075301b5c6d25bb2014cc344b1d18bd Mon Sep 17 00:00:00 2001
From: VOrishirne <vorishirne@gmail.com>
Date: Tue, 7 Feb 2023 16:40:59 +0530
Subject: [PATCH] extend viper's pflag binding to stringToInt pflag

---
 viper.go      | 29 ++++++++++++++++++++++++++++-
 viper_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/viper.go b/viper.go
index 7eac4b7..d519e17 100644
--- a/viper.go
+++ b/viper.go
@@ -1289,7 +1289,8 @@ func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} {
 			return cast.ToDurationSlice(slice)
 		case "stringToString":
 			return stringToStringConv(flag.ValueString())
-
+		case "stringToInt":
+			return stringToIntConv(flag.ValueString())
 		default:
 			return flag.ValueString()
 		}
@@ -1370,6 +1371,8 @@ func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} {
 				return cast.ToIntSlice(res)
 			case "stringToString":
 				return stringToStringConv(flag.ValueString())
+			case "stringToInt":
+				return stringToIntConv(flag.ValueString())
 			case "durationSlice":
 				s := strings.TrimPrefix(flag.ValueString(), "[")
 				s = strings.TrimSuffix(s, "]")
@@ -1418,6 +1421,30 @@ func stringToStringConv(val string) interface{} {
 	return out
 }
 
+// mostly copied from pflag's implementation of this operation here https://github.com/spf13/pflag/blob/d5e0c0615acee7028e1e2740a11102313be88de1/string_to_int.go#L68
+// alterations are: errors are swallowed, map[string]interface{} is returned in order to enable cast.ToStringMap
+func stringToIntConv(val string) interface{} {
+	val = strings.Trim(val, "[]")
+	// An empty string would cause an empty map
+	if len(val) == 0 {
+		return map[string]interface{}{}
+	}
+	ss := strings.Split(val, ",")
+	out := make(map[string]interface{}, len(ss))
+	for _, pair := range ss {
+		kv := strings.SplitN(pair, "=", 2)
+		if len(kv) != 2 {
+			return nil
+		}
+		var err error
+		out[kv[0]], err = strconv.Atoi(kv[1])
+		if err != nil {
+			return nil
+		}
+	}
+	return out
+}
+
 // IsSet checks to see if the key has been set in any of the data locations.
 // IsSet is case-insensitive for a key.
 func IsSet(key string) bool { return v.IsSet(key) }
diff --git a/viper_test.go b/viper_test.go
index 8283b5c..b48c950 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -1255,6 +1255,53 @@ func TestBindPFlagStringToString(t *testing.T) {
 	}
 }
 
+func TestBindPFlagStringToInt(t *testing.T) {
+	tests := []struct {
+		Expected map[string]int
+		Value    string
+	}{
+		{map[string]int{"yo": 1, "oh": 21}, "yo=1,oh=21"},
+		{map[string]int{"yo": 100000000, "oh": 0}, "yo=100000000,oh=0"},
+		{map[string]int{}, "yo=2,oh=21.0"},
+		{map[string]int{}, "yo=,oh=20.99"},
+		{map[string]int{}, "yo=,oh="},
+	}
+
+	v := New() // create independent Viper object
+	defaultVal := map[string]int{}
+	v.SetDefault("stringtoint", defaultVal)
+
+	for _, testValue := range tests {
+		flagSet := pflag.NewFlagSet("test", pflag.ContinueOnError)
+		flagSet.StringToInt("stringtoint", testValue.Expected, "test")
+
+		for _, changed := range []bool{true, false} {
+			flagSet.VisitAll(func(f *pflag.Flag) {
+				f.Value.Set(testValue.Value)
+				f.Changed = changed
+			})
+
+			err := v.BindPFlags(flagSet)
+			if err != nil {
+				t.Fatalf("error binding flag set, %v", err)
+			}
+
+			type TestMap struct {
+				StringToInt map[string]int
+			}
+			val := &TestMap{}
+			if err := v.Unmarshal(val); err != nil {
+				t.Fatalf("%+#v cannot unmarshal: %s", testValue.Value, err)
+			}
+			if changed {
+				assert.Equal(t, testValue.Expected, val.StringToInt)
+			} else {
+				assert.Equal(t, defaultVal, val.StringToInt)
+			}
+		}
+	}
+}
+
 func TestBoundCaseSensitivity(t *testing.T) {
 	assert.Equal(t, "brown", Get("eyes"))