diff --git a/toml.go b/toml.go index d323c39..2b138d1 100644 --- a/toml.go +++ b/toml.go @@ -122,6 +122,89 @@ func (t *Tree) GetPath(keys []string) interface{} { } } +// GetArray returns the value at key in the Tree. +// It returns []string, []int64, etc type if key has homogeneous lists +// Key is a dot-separated path (e.g. a.b.c) without single/double quoted strings. +// Returns nil if the path does not exist in the tree. +// If keys is of length zero, the current tree is returned. +func (t *Tree) GetArray(key string) interface{} { + if key == "" { + return t + } + return t.GetArrayPath(strings.Split(key, ".")) +} + +// GetArrayPath returns the element in the tree indicated by 'keys'. +// If keys is of length zero, the current tree is returned. +func (t *Tree) GetArrayPath(keys []string) interface{} { + if len(keys) == 0 { + return t + } + subtree := t + for _, intermediateKey := range keys[:len(keys)-1] { + value, exists := subtree.values[intermediateKey] + if !exists { + return nil + } + switch node := value.(type) { + case *Tree: + subtree = node + case []*Tree: + // go to most recent element + if len(node) == 0 { + return nil + } + subtree = node[len(node)-1] + default: + return nil // cannot navigate through other node types + } + } + // branch based on final node type + switch node := subtree.values[keys[len(keys)-1]].(type) { + case *tomlValue: + switch n := node.value.(type) { + case []interface{}: + return getArray(n) + default: + return node.value + } + default: + return node + } +} + +// if homogeneous array, then return slice type object over []interface{} +func getArray(n []interface{}) interface{} { + var s []string + var i64 []int64 + var f64 []float64 + var bl []bool + for _, value := range n { + switch v := value.(type) { + case string: + s = append(s, v) + case int64: + i64 = append(i64, v) + case float64: + f64 = append(f64, v) + case bool: + bl = append(bl, v) + default: + return n + } + } + if len(s) == len(n) { + return s + } else if len(i64) == len(n) { + return i64 + } else if len(f64) == len(n) { + return f64 + } else if len(bl) == len(n) { + return bl + } + return n +} + // GetPosition returns the position of the given key. func (t *Tree) GetPosition(key string) Position { if key == "" { diff --git a/toml_test.go b/toml_test.go index 81c3954..0c7b6d3 100644 --- a/toml_test.go +++ b/toml_test.go @@ -3,6 +3,7 @@ package toml import ( + "reflect" "testing" ) @@ -39,6 +40,41 @@ func TestTomlGet(t *testing.T) { } } +func TestTomlGetArray(t *testing.T) { + tree, _ := Load(` + [test] + key = ["one", "two"] + key2 = [true, false, false] + key3 = [1.5,2.5] + `) + + if tree.GetArray("") != tree { + t.Errorf("GetArray should return the tree itself when given an empty path") + } + + expect := []string{"one", "two"} + actual := tree.GetArray("test.key").([]string) + if !reflect.DeepEqual(actual, expect) { + t.Errorf("GetArray should return the []string value") + } + + expect2 := []bool{true, false, false} + actual2 := tree.GetArray("test.key2").([]bool) + if !reflect.DeepEqual(actual2, expect2) { + t.Errorf("GetArray should return the []bool value") + } + + expect3 := []float64{1.5, 2.5} + actual3 := tree.GetArray("test.key3").([]float64) + if !reflect.DeepEqual(actual3, expect3) { + t.Errorf("GetArray should return the []float64 value") + } + + if tree.GetArray(`\`) != nil { + t.Errorf("should return nil when the key is malformed") + } +} + func TestTomlGetDefault(t *testing.T) { tree, _ := Load(` [test] @@ -148,6 +184,51 @@ func TestTomlGetPath(t *testing.T) { } } +func TestTomlGetArrayPath(t *testing.T) { + for idx, item := range []struct { + Name string + Path []string + Make func() (tree *Tree, expected interface{}) + }{ + { + Name: "empty", + Path: []string{}, + Make: func() (tree *Tree, expected interface{}) { + tree = newTree() + expected = tree + return + }, + }, + { + Name: "int64", + Path: []string{"a"}, + Make: func() (tree *Tree, expected interface{}) { + var err error + tree, err = Load(`a = [1,2,3]`) + if err != nil { + panic(err) + } + expected = []int64{1, 2, 3} + return + }, + }, + } { + t.Run(item.Name, func(t *testing.T) { + tree, expected := item.Make() + result := tree.GetArrayPath(item.Path) + if !reflect.DeepEqual(result, expected) { + t.Errorf("GetArrayPath[%d] %v - expected %#v, got %#v instead.", idx, item.Path, expected, result) + } + }) + } + + tree, _ := Load("[foo.bar]\na=1\nb=2\n[baz.foo]\na=3\nb=4\n[gorf.foo]\na=5\nb=6") + if tree.GetArrayPath([]string{"whatever"}) != nil { + t.Error("GetArrayPath should return nil when the key does not exist") + } + +} + func TestTomlFromMap(t *testing.T) { simpleMap := map[string]interface{}{"hello": 42} tree, err := TreeFromMap(simpleMap)