From e5459cc690270053fa7d374d4b72a4398f38229e Mon Sep 17 00:00:00 2001 From: Harley Laue Date: Mon, 23 Apr 2018 14:27:51 -0700 Subject: [PATCH] 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")