Skip to content

Commit 2ee1631

Browse files
vorishirnesagikazarmark
authored andcommitted
extend viper's pflag binding to stringToInt pflag
1 parent 8652057 commit 2ee1631

File tree

2 files changed

+75
-1
lines changed

2 files changed

+75
-1
lines changed

viper.go

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1289,7 +1289,8 @@ func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} {
12891289
return cast.ToDurationSlice(slice)
12901290
case "stringToString":
12911291
return stringToStringConv(flag.ValueString())
1292-
1292+
case "stringToInt":
1293+
return stringToIntConv(flag.ValueString())
12931294
default:
12941295
return flag.ValueString()
12951296
}
@@ -1370,6 +1371,8 @@ func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} {
13701371
return cast.ToIntSlice(res)
13711372
case "stringToString":
13721373
return stringToStringConv(flag.ValueString())
1374+
case "stringToInt":
1375+
return stringToIntConv(flag.ValueString())
13731376
case "durationSlice":
13741377
s := strings.TrimPrefix(flag.ValueString(), "[")
13751378
s = strings.TrimSuffix(s, "]")
@@ -1418,6 +1421,30 @@ func stringToStringConv(val string) interface{} {
14181421
return out
14191422
}
14201423

1424+
// mostly copied from pflag's implementation of this operation here https://github.com/spf13/pflag/blob/d5e0c0615acee7028e1e2740a11102313be88de1/string_to_int.go#L68
1425+
// alterations are: errors are swallowed, map[string]interface{} is returned in order to enable cast.ToStringMap
1426+
func stringToIntConv(val string) interface{} {
1427+
val = strings.Trim(val, "[]")
1428+
// An empty string would cause an empty map
1429+
if len(val) == 0 {
1430+
return map[string]interface{}{}
1431+
}
1432+
ss := strings.Split(val, ",")
1433+
out := make(map[string]interface{}, len(ss))
1434+
for _, pair := range ss {
1435+
kv := strings.SplitN(pair, "=", 2)
1436+
if len(kv) != 2 {
1437+
return nil
1438+
}
1439+
var err error
1440+
out[kv[0]], err = strconv.Atoi(kv[1])
1441+
if err != nil {
1442+
return nil
1443+
}
1444+
}
1445+
return out
1446+
}
1447+
14211448
// IsSet checks to see if the key has been set in any of the data locations.
14221449
// IsSet is case-insensitive for a key.
14231450
func IsSet(key string) bool { return v.IsSet(key) }

viper_test.go

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1255,6 +1255,53 @@ func TestBindPFlagStringToString(t *testing.T) {
12551255
}
12561256
}
12571257

1258+
func TestBindPFlagStringToInt(t *testing.T) {
1259+
tests := []struct {
1260+
Expected map[string]int
1261+
Value string
1262+
}{
1263+
{map[string]int{"yo": 1, "oh": 21}, "yo=1,oh=21"},
1264+
{map[string]int{"yo": 100000000, "oh": 0}, "yo=100000000,oh=0"},
1265+
{map[string]int{}, "yo=2,oh=21.0"},
1266+
{map[string]int{}, "yo=,oh=20.99"},
1267+
{map[string]int{}, "yo=,oh="},
1268+
}
1269+
1270+
v := New() // create independent Viper object
1271+
defaultVal := map[string]int{}
1272+
v.SetDefault("stringtoint", defaultVal)
1273+
1274+
for _, testValue := range tests {
1275+
flagSet := pflag.NewFlagSet("test", pflag.ContinueOnError)
1276+
flagSet.StringToInt("stringtoint", testValue.Expected, "test")
1277+
1278+
for _, changed := range []bool{true, false} {
1279+
flagSet.VisitAll(func(f *pflag.Flag) {
1280+
f.Value.Set(testValue.Value)
1281+
f.Changed = changed
1282+
})
1283+
1284+
err := v.BindPFlags(flagSet)
1285+
if err != nil {
1286+
t.Fatalf("error binding flag set, %v", err)
1287+
}
1288+
1289+
type TestMap struct {
1290+
StringToInt map[string]int
1291+
}
1292+
val := &TestMap{}
1293+
if err := v.Unmarshal(val); err != nil {
1294+
t.Fatalf("%+#v cannot unmarshal: %s", testValue.Value, err)
1295+
}
1296+
if changed {
1297+
assert.Equal(t, testValue.Expected, val.StringToInt)
1298+
} else {
1299+
assert.Equal(t, defaultVal, val.StringToInt)
1300+
}
1301+
}
1302+
}
1303+
}
1304+
12581305
func TestBoundCaseSensitivity(t *testing.T) {
12591306
assert.Equal(t, "brown", Get("eyes"))
12601307

0 commit comments

Comments
 (0)