Add map support to scratch

2 files changed, 54 insertions, 0 deletions

diff --git a/hugolib/scratch.go b/hugolib/scratch.go
index 0f5c4b484..16d1c1440 100644
--- a/hugolib/scratch.go
+++ b/hugolib/scratch.go
@@ -15,6 +15,7 @@ package hugolib
 
 import (
 	"github.com/spf13/hugo/helpers"
+	"sort"
 )
 
 // Scratch is a writable context used for stateful operations in Page/Node rendering.
@@ -52,6 +53,41 @@ func (c *Scratch) Get(key string) interface{} {
 	return c.values[key]
 }
 
+// SetInMap stores a value to a map with the given key in the Node context.
+// This map can later be retrieved with GetSortedMapValues.
+func (c *Scratch) SetInMap(key string, mapKey string, value interface{}) string {
+	_, found := c.values[key]
+	if !found {
+		c.values[key] = make(map[string]interface{})
+	}
+
+	c.values[key].(map[string]interface{})[mapKey] = value
+	return ""
+}
+
+// GetSortedMapValues returns a sorted map previously filled with SetInMap
+func (c *Scratch) GetSortedMapValues(key string) interface{} {
+	if c.values[key] == nil {
+		return nil
+	}
+
+	unsortedMap := c.values[key].(map[string]interface{})
+
+	var keys []string
+	for mapKey, _ := range unsortedMap {
+		keys = append(keys, mapKey)
+	}
+
+	sort.Strings(keys)
+
+	sortedArray := make([]interface{}, len(unsortedMap))
+	for i, mapKey := range keys {
+		sortedArray[i] = unsortedMap[mapKey]
+	}
+
+	return sortedArray
+}
+
 func newScratch() *Scratch {
 	return &Scratch{values: make(map[string]interface{})}
 }
diff --git a/hugolib/scratch_test.go b/hugolib/scratch_test.go
index c736f5a99..b49d8982b 100644
--- a/hugolib/scratch_test.go
+++ b/hugolib/scratch_test.go
@@ -47,3 +47,21 @@ func TestScratchGet(t *testing.T) {
 		t.Errorf("Should not return anything, but got %v", nothing)
 	}
 }
+
+func TestScratchSetInMap(t *testing.T) {
+	scratch := newScratch()
+	scratch.SetInMap("key", "lux", "Lux")
+	scratch.SetInMap("key", "abc", "Abc")
+	scratch.SetInMap("key", "zyx", "Zyx")
+	scratch.SetInMap("key", "abc", "Abc (updated)")
+	scratch.SetInMap("key", "def", "Def")
+	assert.Equal(t, []interface{}{0: "Abc (updated)", 1: "Def", 2: "Lux", 3: "Zyx"}, scratch.GetSortedMapValues("key"))
+}
+
+func TestScratchGetSortedMapValues(t *testing.T) {
+	scratch := newScratch()
+	nothing := scratch.GetSortedMapValues("nothing")
+	if nothing != nil {
+		t.Errorf("Should not return anything, but got %v", nothing)
+	}
+}