From bcacc71a18be4b36b52baf6d10a95513f94dc7b2 Mon Sep 17 00:00:00 2001
From: RiyaJohn <riyajohn29@rediffmail.com>
Date: Mon, 18 May 2020 15:26:15 +0530
Subject: [PATCH] feat: add GetArray() with testcases

---
 toml.go      | 66 ++++++++++++++++++++++++++++++++++++++++++++++++----
 toml_test.go | 48 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+), 5 deletions(-)

diff --git a/toml.go b/toml.go
index 7a580c9..fdc74d8 100644
--- a/toml.go
+++ b/toml.go
@@ -90,6 +90,52 @@ func (t *Tree) Get(key string) interface{} {
 // GetPath returns the element in the tree indicated by 'keys'.
 // If keys is of length zero, the current tree is returned.
 func (t *Tree) GetPath(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:
+		return node.value
+	default:
+		return node
+	}
+}
+
+// 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
 	}
@@ -126,18 +172,22 @@ func (t *Tree) GetPath(keys []string) interface{} {
 	}
 }
 
-// if homogeneous array, then return array type obj. as resp. over []interface{}
+// if homogeneous array, then return slice type object over []interface{}
 func getArray(n []interface{}) interface{} {
-	var str []string
+	var s []string
+	var b []byte
 	var i32 []int32
 	var i64 []int64
 	var i []int
 	var f32 []float32
 	var f64 []float64
+	var bl []bool
 	for _, value := range n {
 		switch v := value.(type) {
 		case string:
-			str = append(str, v)
+			s = append(s, v)
+		case byte:
+			b = append(b, v)
 		case int32:
 			i32 = append(i32, v)
 		case int64:
@@ -148,12 +198,16 @@ func getArray(n []interface{}) interface{} {
 			f32 = append(f32, v)
 		case float64:
 			f64 = append(f64, v)
+		case bool:
+			bl = append(bl, v)
 		default:
 			return n
 		}
 	}
-	if len(str) == len(n) {
-		return str
+	if len(s) == len(n) {
+		return s
+	} else if len(b) == len(n) {
+		return b
 	} else if len(i32) == len(n) {
 		return i32
 	} else if len(i64) == len(n) {
@@ -164,6 +218,8 @@ func getArray(n []interface{}) interface{} {
 		return f32
 	} else if len(f64) == len(n) {
 		return f64
+	} else if len(bl) == len(n) {
+		return bl
 	}
 	return n
 }
diff --git a/toml_test.go b/toml_test.go
index 81c3954..ce3cf1c 100644
--- a/toml_test.go
+++ b/toml_test.go
@@ -3,6 +3,7 @@
 package toml
 
 import (
+	"reflect"
 	"testing"
 )
 
@@ -39,6 +40,27 @@ func TestTomlGet(t *testing.T) {
 	}
 }
 
+func TestTomlGetArray(t *testing.T) {
+	tree, _ := Load(`
+		[test]
+		key = ["one", "two"]
+	`)
+
+	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")
+	}
+
+	if tree.GetArray(`\`) != nil {
+		t.Errorf("should return nil when the key is malformed")
+	}
+}
+
 func TestTomlGetDefault(t *testing.T) {
 	tree, _ := Load(`
 		[test]
@@ -148,6 +170,32 @@ func TestTomlGetPath(t *testing.T) {
 	}
 }
 
+func TestTomlGetArrayPath(t *testing.T) {
+	node := newTree()
+	//TODO: set other node data
+
+	for idx, item := range []struct {
+		Path     []string
+		Expected *Tree
+	}{
+		{ // empty path test
+			[]string{},
+			node,
+		},
+	} {
+		result := node.GetArrayPath(item.Path)
+		if result != item.Expected {
+			t.Errorf("GetArrayPath[%d] %v - expected %v, got %v instead.", idx, item.Path, item.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)