From a7a5948b15ddd0aff0da29a4c4b0011c2968ccdb Mon Sep 17 00:00:00 2001 From: Harley Laue Date: Mon, 23 Apr 2018 10:24:03 -0700 Subject: [PATCH 1/3] Check for nil before binding pflag(s) * When passing nil to BindPFlag or BindPFlags, the value is set to a struct and passed as an interface. That struct never checks for the flag(set) being nil. Thus, it makes sense to check before it's set to the struct. * fixes #422 --- viper.go | 6 ++++++ viper_test.go | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/viper.go b/viper.go index e9966ba..c166e9f 100644 --- a/viper.go +++ b/viper.go @@ -811,6 +811,9 @@ func (v *Viper) UnmarshalExact(rawVal interface{}) error { // name as the config key. func BindPFlags(flags *pflag.FlagSet) error { return v.BindPFlags(flags) } func (v *Viper) BindPFlags(flags *pflag.FlagSet) error { + if flags == nil { + return fmt.Errorf("FlagSet cannot be nil") + } return v.BindFlagValues(pflagValueSet{flags}) } @@ -822,6 +825,9 @@ func (v *Viper) BindPFlags(flags *pflag.FlagSet) error { // func BindPFlag(key string, flag *pflag.Flag) error { return v.BindPFlag(key, flag) } func (v *Viper) BindPFlag(key string, flag *pflag.Flag) error { + if flag == nil { + return fmt.Errorf("flag for %q is nil", key) + } return v.BindFlagValue(key, pflagValue{flag}) } diff --git a/viper_test.go b/viper_test.go index c93480e..443345e 100644 --- a/viper_test.go +++ b/viper_test.go @@ -577,6 +577,14 @@ func TestBindPFlagsStringSlice(t *testing.T) { } } +func TestBindPFlagsNil(t *testing.T) { + v := New() + err := v.BindPFlags(nil) + if err == nil { + t.Fatalf("expected error when passing nil to BindPFlags") + } +} + func TestBindPFlag(t *testing.T) { var testString = "testing" var testValue = newStringValue(testString, &testString) @@ -598,6 +606,14 @@ func TestBindPFlag(t *testing.T) { } +func TestBindPFlagNil(t *testing.T) { + v := New() + err := v.BindPFlag("any", nil) + if err == nil { + t.Fatalf("expected error when passing nil to BindPFlag") + } +} + func TestBoundCaseSensitivity(t *testing.T) { assert.Equal(t, "brown", Get("eyes")) From b7a4909ef7e105ba69078910b7e3adee468ae549 Mon Sep 17 00:00:00 2001 From: Harley Laue Date: Mon, 23 Apr 2018 11:06:26 -0700 Subject: [PATCH 2/3] Allow getting Sub from a prefix of pflags key(s) * It seems common enough that flags may be defined from PFlags be retrieved as a sub Viper similar to getting it from Yaml or other configuration structures. This allows flags to be bound one-to-one with how a config file may be structured. * This is partially based on #331 for the find boolean flag * fixes #368 --- flags_test.go | 2 +- viper.go | 25 ++++++++++++++++++++----- viper_test.go | 34 +++++++++++++++++++++++++++++----- 3 files changed, 50 insertions(+), 11 deletions(-) diff --git a/flags_test.go b/flags_test.go index 0b976b6..1dd7b0d 100644 --- a/flags_test.go +++ b/flags_test.go @@ -45,7 +45,7 @@ func TestBindFlagValueSet(t *testing.T) { func TestBindFlagValue(t *testing.T) { var testString = "testing" - var testValue = newStringValue(testString, &testString) + var testValue = newStringValue(testString) flag := &pflag.Flag{ Name: "testflag", diff --git a/viper.go b/viper.go index c166e9f..104a657 100644 --- a/viper.go +++ b/viper.go @@ -612,7 +612,7 @@ func GetViper() *Viper { func Get(key string) interface{} { return v.Get(key) } func (v *Viper) Get(key string) interface{} { lcaseKey := strings.ToLower(key) - val := v.find(lcaseKey) + val := v.find(lcaseKey, true) if val == nil { return nil } @@ -885,9 +885,12 @@ func (v *Viper) BindEnv(input ...string) error { // Given a key, find the value. // Viper will check in the following order: // flag, env, config file, key/value store, default. -// Viper will check to see if an alias exists first. -// Note: this assumes a lower-cased key given. -func (v *Viper) find(lcaseKey string) interface{} { +// Viper will then check in the following order: +// flag, env, config file, key/value store. +// Lastly, if no value was found and flagDefault is true, and if the key +// corresponds to a flag, the flag's default value is returned. +// +func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} { var ( val interface{} @@ -1003,6 +1006,18 @@ func (v *Viper) find(lcaseKey string) interface{} { } // last item, no need to check shadowing + // it could also be a key prefix, search for that prefix to get the values from + // pflags that match it + sub := make(map[string]interface{}) + for key, val := range v.pflags { + if flagDefault && strings.HasPrefix(key, lcaseKey) { + sub[strings.TrimPrefix(key, lcaseKey+".")] = val.ValueString() + } + } + if len(sub) != 0 { + return sub + } + return nil } @@ -1020,7 +1035,7 @@ func readAsCSV(val string) ([]string, error) { func IsSet(key string) bool { return v.IsSet(key) } func (v *Viper) IsSet(key string) bool { lcaseKey := strings.ToLower(key) - val := v.find(lcaseKey) + val := v.find(lcaseKey, false) return val != nil } diff --git a/viper_test.go b/viper_test.go index 443345e..84326f5 100644 --- a/viper_test.go +++ b/viper_test.go @@ -224,9 +224,8 @@ func initDirs(t *testing.T) (string, string, func()) { //stubs for PFlag Values type stringValue string -func newStringValue(val string, p *string) *stringValue { - *p = val - return (*stringValue)(p) +func newStringValue(val string) *stringValue { + return (*stringValue)(&val) } func (s *stringValue) Set(val string) error { @@ -587,7 +586,7 @@ func TestBindPFlagsNil(t *testing.T) { func TestBindPFlag(t *testing.T) { var testString = "testing" - var testValue = newStringValue(testString, &testString) + var testValue = newStringValue(testString) flag := &pflag.Flag{ Name: "testflag", @@ -623,7 +622,7 @@ func TestBoundCaseSensitivity(t *testing.T) { assert.Equal(t, "blue", Get("eyes")) var testString = "green" - var testValue = newStringValue(testString, &testString) + var testValue = newStringValue(testString) flag := &pflag.Flag{ Name: "eyeballs", @@ -864,6 +863,31 @@ func TestSub(t *testing.T) { assert.Equal(t, (*Viper)(nil), subv) } +func TestSubPflags(t *testing.T) { + v := New() + + // same as yamlExample, without hobbies + v.BindPFlag("name", &pflag.Flag{Value: newStringValue("steve"), Changed: true}) + v.BindPFlag("clothing.jacket", &pflag.Flag{Value: newStringValue("leather"), Changed: true}) + v.BindPFlag("clothing.trousers", &pflag.Flag{Value: newStringValue("denim"), Changed: true}) + v.BindPFlag("clothing.pants.size", &pflag.Flag{Value: newStringValue("large"), Changed: true}) + v.BindPFlag("age", &pflag.Flag{Value: newStringValue("35"), Changed: true}) + v.BindPFlag("eyes", &pflag.Flag{Value: newStringValue("brown"), Changed: true}) + v.BindPFlag("beard", &pflag.Flag{Value: newStringValue("yes"), Changed: true}) + + subv := v.Sub("clothing") + assert.Equal(t, v.Get("clothing.pants.size"), subv.Get("pants.size")) + + subv = v.Sub("clothing.pants") + assert.Equal(t, v.Get("clothing.pants.size"), subv.Get("size")) + + subv = v.Sub("clothing.pants.size") + assert.Equal(t, (*Viper)(nil), subv) + + subv = v.Sub("missing.key") + assert.Equal(t, (*Viper)(nil), subv) +} + var hclWriteExpected = []byte(`"foos" = { "foo" = { "key" = 1 From e5459cc690270053fa7d374d4b72a4398f38229e Mon Sep 17 00:00:00 2001 From: Harley Laue Date: Mon, 23 Apr 2018 14:27:51 -0700 Subject: [PATCH 3/3] Change implementation of find sub-tree to work with deep trees * The previous implementation had an issue with nested structures deeper than one level. It would copy keys over instead of the expected object structure. In the example in this commit, UnmarshalKey("clothing") was returning "pants.size":"35" instead of "pants"{"size":"35"} --- viper.go | 12 +++++++++--- viper_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/viper.go b/viper.go index 104a657..abd9c9e 100644 --- a/viper.go +++ b/viper.go @@ -1009,9 +1009,15 @@ func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} { // it could also be a key prefix, search for that prefix to get the values from // pflags that match it sub := make(map[string]interface{}) - for key, val := range v.pflags { - if flagDefault && strings.HasPrefix(key, lcaseKey) { - sub[strings.TrimPrefix(key, lcaseKey+".")] = val.ValueString() + for _, key := range v.AllKeys() { + if strings.HasPrefix(key, lcaseKey) { + value := v.Get(key) + keypath := strings.Split(lcaseKey, v.keyDelim) + path := strings.Split(key, v.keyDelim)[len(keypath)-1:] + lastKey := strings.ToLower(path[len(path)-1]) + deepestMap := deepSearch(sub, path[1:len(path)-1]) + // set innermost value + deepestMap[lastKey] = value } } if len(sub) != 0 { diff --git a/viper_test.go b/viper_test.go index 84326f5..f886358 100644 --- a/viper_test.go +++ b/viper_test.go @@ -875,7 +875,49 @@ func TestSubPflags(t *testing.T) { v.BindPFlag("eyes", &pflag.Flag{Value: newStringValue("brown"), Changed: true}) v.BindPFlag("beard", &pflag.Flag{Value: newStringValue("yes"), Changed: true}) + type pants struct { + Size string + } + + type clothing struct { + Jacket string + Trousers string + Pants pants + } + + type cfg struct { + Name string + Clothing clothing + Age int + Eyes string + Beard bool + } + + var c cfg + v.Unmarshal(&c) + assert.Equal(t, v.Get("name"), c.Name) + assert.Equal(t, v.Get("clothing.jacket"), c.Clothing.Jacket) + assert.Equal(t, v.Get("clothing.trousers"), c.Clothing.Trousers) + assert.Equal(t, v.Get("clothing.pants.size"), c.Clothing.Pants.Size) + assert.Equal(t, v.GetInt("age"), c.Age) + assert.Equal(t, v.Get("eyes"), c.Eyes) + assert.Equal(t, v.GetBool("beard"), c.Beard) + + var cloth clothing + v.UnmarshalKey("clothing", &cloth) + assert.Equal(t, c.Clothing, cloth) + + var p pants + v.UnmarshalKey("clothing.pants", &p) + assert.Equal(t, c.Clothing.Pants, p) + + var size string + v.UnmarshalKey("clothing.pants.size", &size) + assert.Equal(t, c.Clothing.Pants.Size, size) + subv := v.Sub("clothing") + assert.Equal(t, v.Get("clothing.jacket"), subv.Get("jacket")) + assert.Equal(t, v.Get("clothing.trousers"), subv.Get("trousers")) assert.Equal(t, v.Get("clothing.pants.size"), subv.Get("pants.size")) subv = v.Sub("clothing.pants")