Simplify, document and add more tests to var replacement.

api/filters/refvar/doc.go

@@ -1,3 +1,3 @@
 // Package refvar contains a kio.Filter implementation of the kustomize
-// refvar transformer.
+// refvar transformer (find and replace $(FOO) style variables in strings).
 package refvar

api/filters/refvar/expand.go

@@ -1,12 +1,12 @@
 // Copyright 2019 The Kubernetes Authors.
 // SPDX-License-Identifier: Apache-2.0
 
-// Package expansion provides functions find and replace $(FOO) style variables in strings.
 package refvar
 
 import (
-	"bytes"
 	"fmt"
+	"log"
+	"strings"
 )
 
 const (
@@ -17,38 +17,64 @@ const (
 
 // syntaxWrap returns the input string wrapped by the expansion syntax.
 func syntaxWrap(input string) string {
-	return string(operator) + string(referenceOpener) + input + string(referenceCloser)
+	var sb strings.Builder
+	sb.WriteByte(operator)
+	sb.WriteByte(referenceOpener)
+	sb.WriteString(input)
+	sb.WriteByte(referenceCloser)
+	return sb.String()
 }
 
-// MappingFuncFor returns a mapping function for use with Expand that
-// implements the expansion semantics defined in the expansion spec; it
-// returns the input string wrapped in the expansion syntax if no mapping
-// for the input is found.
-func MappingFuncFor(
-	counts map[string]int,
-	context ...map[string]interface{}) func(string) interface{} {
-	return func(input string) interface{} {
-		for _, vars := range context {
-			val, ok := vars[input]
-			if ok {
-				counts[input]++
-				switch typedV := val.(type) {
-				case string, int64, float64, bool:
-					return typedV
-				default:
-					return syntaxWrap(input)
-				}
+// MappingFunc maps a string to anything.
+type MappingFunc func(string) interface{}
+
+// MakePrimitiveReplacer returns a MappingFunc that uses a map to do
+// replacements, and a histogram to count map hits.
+//
+// Func behavior:
+//
+// If the input key is NOT found in the map, the key is wrapped up as
+// as a variable declaration string and returned, e.g. key FOO becomes $(FOO).
+// This string is presumably put back where it was found, and might get replaced
+// later.
+//
+// If the key is found in the map, the value is returned if it is a primitive
+// type (string, bool, number), and the hit is counted.
+//
+// If it's not a primitive type (e.g. a map, struct, func, etc.) then this
+// function doesn't know what to do with it and it returns the key wrapped up
+// again as if it had not been replaced.  This should probably be an error.
+func MakePrimitiveReplacer(
+	counts map[string]int, someMap map[string]interface{}) MappingFunc {
+	return func(key string) interface{} {
+		if value, ok := someMap[key]; ok {
+			switch typedV := value.(type) {
+			case string, int, int32, int64, float32, float64, bool:
+				counts[key]++
+				return typedV
+			default:
+				// If the value is some complicated type (e.g. a map or struct),
+				// this function doesn't know how to jam it into a string,
+				// so just pretend it was a cache miss.
+				// Likely this should be an error instead of a silent failure,
+				// since the programmer passed an impossible value.
+				log.Printf(
+					"MakePrimitiveReplacer: bad replacement type=%T val=%v",
+					typedV, typedV)
+				return syntaxWrap(key)
 			}
 		}
-		return syntaxWrap(input)
+		// If unable to return the mapped variable, return it
+		// as it was found, and a later mapping might be able to
+		// replace it.
+		return syntaxWrap(key)
 	}
 }
 
-// Expand replaces variable references in the input string according to
-// the expansion spec using the given mapping function to resolve the
-// values of variables.
-func Expand(input string, mapping func(string) interface{}) interface{} {
-	var buf bytes.Buffer
+// DoReplacements replaces variable references in the input string
+// using the mapping function.
+func DoReplacements(input string, mapping MappingFunc) interface{} {
+	var buf strings.Builder
 	checkpoint := 0
 	for cursor := 0; cursor < len(input); cursor++ {
 		if input[cursor] == operator && cursor+1 < len(input) {

api/filters/refvar/expand_test.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	. "sigs.k8s.io/kustomize/api/filters/refvar"
 )
 
@@ -15,6 +16,48 @@ type expected struct {
 	edited string
 }
 
+func TestPrimitiveReplacer(t *testing.T) {
+	varCounts := make(map[string]int)
+	f := MakePrimitiveReplacer(
+		varCounts,
+		map[string]interface{}{
+			"FOO":        "bar",
+			"ZOO":        "$(FOO)-1",
+			"BLU":        "$(ZOO)-2",
+			"EIGHT":      8,
+			"PI":         3.14159,
+			"ZINT":       "$(INT)",
+			"BOOL":       "true",
+			"HUGENUMBER": int64(9223372036854775807),
+			"CRAZYMAP":   map[string]int{"crazy": 200},
+			"ZBOOL":      "$(BOOL)",
+		})
+	assert.Equal(t, "$()", f(""))
+	assert.Equal(t, "$(      )", f("      "))
+	assert.Equal(t, "$(florida)", f("florida"))
+	assert.Equal(t, "$(0)", f("0"))
+	assert.Equal(t, "bar", f("FOO"))
+	assert.Equal(t, "bar", f("FOO"))
+	assert.Equal(t, "bar", f("FOO"))
+	assert.Equal(t, 8, f("EIGHT"))
+	assert.Equal(t, 8, f("EIGHT"))
+	assert.Equal(t, 3.14159, f("PI"))
+	assert.Equal(t, "true", f("BOOL"))
+	assert.Equal(t, int64(9223372036854775807), f("HUGENUMBER"))
+	assert.Equal(t, "$(FOO)-1", f("ZOO"))
+	assert.Equal(t, "$(CRAZYMAP)", f("CRAZYMAP"))
+	assert.Equal(t,
+		map[string]int{
+			"FOO":        3,
+			"EIGHT":      2,
+			"BOOL":       1,
+			"PI":         1,
+			"ZOO":        1,
+			"HUGENUMBER": 1,
+		},
+		varCounts)
+}
+
 func TestMapReference(t *testing.T) {
 	type env struct {
 		Name  string
@@ -51,7 +94,7 @@ func TestMapReference(t *testing.T) {
 		},
 	}
 
-	declaredEnv := map[string]interface{}{
+	varMap := map[string]interface{}{
 		"FOO":   "bar",
 		"ZOO":   "$(FOO)-1",
 		"BLU":   "$(ZOO)-2",
@@ -61,11 +104,11 @@ func TestMapReference(t *testing.T) {
 		"ZBOOL": "$(BOOL)",
 	}
 
-	counts := make(map[string]int)
-	mapping := MappingFuncFor(counts, declaredEnv)
-
+	varCounts := make(map[string]int)
 	for _, env := range envs {
-		declaredEnv[env.Name] = Expand(fmt.Sprintf("%v", env.Value), mapping)
+		varMap[env.Name] = DoReplacements(
+			fmt.Sprintf("%v", env.Value),
+			MakePrimitiveReplacer(varCounts, varMap))
 	}
 
 	expectedEnv := map[string]expected{
@@ -79,45 +122,20 @@ func TestMapReference(t *testing.T) {
 	}
 
 	for k, v := range expectedEnv {
-		if e, a := v, declaredEnv[k]; e.edited != a || e.count != counts[k] {
+		if e, a := v, varMap[k]; e.edited != a || e.count != varCounts[k] {
 			t.Errorf("Expected %v count=%d, got %v count=%d",
-				e.edited, e.count, a, counts[k])
+				e.edited, e.count, a, varCounts[k])
 		} else {
-			delete(declaredEnv, k)
+			delete(varMap, k)
 		}
 	}
 
-	if len(declaredEnv) != 0 {
-		t.Errorf("Unexpected keys in declared env: %v", declaredEnv)
+	if len(varMap) != 0 {
+		t.Errorf("Unexpected keys in declared env: %v", varMap)
 	}
 }
 
 func TestMapping(t *testing.T) {
-	context := map[string]interface{}{
-		"VAR_A":     "A",
-		"VAR_B":     "B",
-		"VAR_C":     "C",
-		"VAR_REF":   "$(VAR_A)",
-		"VAR_EMPTY": "",
-	}
-	doExpansionTest(t, context)
-}
-
-func TestMappingDual(t *testing.T) {
-	context := map[string]interface{}{
-		"VAR_A":     "A",
-		"VAR_EMPTY": "",
-	}
-	context2 := map[string]interface{}{
-		"VAR_B":   "B",
-		"VAR_C":   "C",
-		"VAR_REF": "$(VAR_A)",
-	}
-
-	doExpansionTest(t, context, context2)
-}
-
-func doExpansionTest(t *testing.T, context ...map[string]interface{}) {
 	cases := []struct {
 		name     string
 		input    string
@@ -333,11 +351,17 @@ func doExpansionTest(t *testing.T, context ...map[string]interface{}) {
 			expected: "\n",
 		},
 	}
-
 	for _, tc := range cases {
 		counts := make(map[string]int)
-		mapping := MappingFuncFor(counts, context...)
-		expanded := Expand(fmt.Sprintf("%v", tc.input), mapping)
+		expanded := DoReplacements(
+			fmt.Sprintf("%v", tc.input),
+			MakePrimitiveReplacer(counts, map[string]interface{}{
+				"VAR_A":     "A",
+				"VAR_B":     "B",
+				"VAR_C":     "C",
+				"VAR_REF":   "$(VAR_A)",
+				"VAR_EMPTY": "",
+			}))
 		if e, a := tc.expected, expanded; e != a {
 			t.Errorf("%v: expected %q, got %q", tc.name, e, a)
 		}
@@ -347,8 +371,7 @@ func doExpansionTest(t *testing.T, context ...map[string]interface{}) {
 		}
 		if len(tc.counts) > 0 {
 			for k, expectedCount := range tc.counts {
-				c, ok := counts[k]
-				if ok {
+				if c, ok := counts[k]; ok {
 					if c != expectedCount {
 						t.Errorf(
 							"%v: k=%s, expected count %d, got %d",

api/filters/refvar/refvar.go

@@ -13,8 +13,8 @@ import (
 // Filter updates $(VAR) style variables with values.
 // The fieldSpecs are the places to look for occurrences of $(VAR).
 type Filter struct {
-	MappingFunc func(string) interface{} `json:"mappingFunc,omitempty" yaml:"mappingFunc,omitempty"`
-	FieldSpec   types.FieldSpec          `json:"fieldSpec,omitempty" yaml:"fieldSpec,omitempty"`
+	MappingFunc MappingFunc     `json:"mappingFunc,omitempty" yaml:"mappingFunc,omitempty"`
+	FieldSpec   types.FieldSpec `json:"fieldSpec,omitempty" yaml:"fieldSpec,omitempty"`
 }
 
 func (f Filter) Filter(nodes []*yaml.RNode) ([]*yaml.RNode, error) {
@@ -67,7 +67,7 @@ func (f Filter) setScalar(node *yaml.RNode) error {
 	if !yaml.IsYNodeString(node.YNode()) {
 		return nil
 	}
-	v := Expand(node.YNode().Value, f.MappingFunc)
+	v := DoReplacements(node.YNode().Value, f.MappingFunc)
 	updateNodeValue(node.YNode(), v)
 	return nil
 }
@@ -83,7 +83,7 @@ func (f Filter) setMap(node *yaml.RNode) error {
 		if !yaml.IsYNodeString(contents[i+1]) {
 			continue
 		}
-		newValue := Expand(contents[i+1].Value, f.MappingFunc)
+		newValue := DoReplacements(contents[i+1].Value, f.MappingFunc)
 		updateNodeValue(contents[i+1], newValue)
 	}
 	return nil
@@ -94,7 +94,7 @@ func (f Filter) setSeq(node *yaml.RNode) error {
 		if !yaml.IsYNodeString(item) {
 			return fmt.Errorf("invalid value type expect a string")
 		}
-		newValue := Expand(item.Value, f.MappingFunc)
+		newValue := DoReplacements(item.Value, f.MappingFunc)
 		updateNodeValue(item, newValue)
 	}
 	return nil