First pass variable support.

This allows blocked scoped variables to be set and retrieved as values using the $ prefix. e.g. foo = 22; bar = $foo Also supports env variables being used as variables and will properly parse to the correct type.
ericsson-ddi · May 22, 2016 · 832bac9 · 832bac9
1 parent 68e8b83
commit 832bac9
Show file tree

Hide file tree

Showing 4 changed files with 139 additions and 2 deletions.
diff --git a/conf/lex.go b/conf/lex.go
@@ -38,6 +38,7 @@ const (
 	itemMapStart
 	itemMapEnd
 	itemCommentStart
+	itemVariable
 )
 
 const (
@@ -178,7 +179,7 @@ func (lx *lexer) errorf(format string, values ...interface{}) stateFn {
 	return nil
 }
 
-// lexTop consumes elements at the top level of TOML data.
+// lexTop consumes elements at the top level of data structure.
 func lexTop(lx *lexer) stateFn {
 	r := lx.next()
 	if isWhitespace(r) || isNL(r) {
@@ -290,7 +291,7 @@ func lexQuotedKey(lx *lexer) stateFn {
 // is not whitespace) has already been consumed.
 func lexKey(lx *lexer) stateFn {
 	r := lx.peek()
-	if isWhitespace(r) || isNL(r) || isKeySeparator(r) {
+	if isWhitespace(r) || isNL(r) || isKeySeparator(r) || r == eof {
 		lx.emit(itemKey)
 		return lexKeyEnd
 	}
@@ -308,6 +309,9 @@ func lexKeyEnd(lx *lexer) stateFn {
 		return lexSkip(lx, lexKeyEnd)
 	case isKeySeparator(r):
 		return lexSkip(lx, lexValue)
+	case r == eof:
+		lx.emit(itemEOF)
+		return nil
 	}
 	// We start the value here
 	lx.backup()
@@ -570,6 +574,15 @@ func (lx *lexer) isBool() bool {
 	return str == "true" || str == "false" || str == "TRUE" || str == "FALSE"
 }
 
+// Check if the unquoted string is a variable reference, starting with $.
+func (lx *lexer) isVariable() bool {
+	if lx.input[lx.start] == '$' {
+		lx.start += 1
+		return true
+	}
+	return false
+}
+
 // lexQuotedString consumes the inner contents of a string. It assumes that the
 // beginning '"' has already been consumed and ignored. It will not interpret any
 // internal contents.
@@ -616,6 +629,8 @@ func lexString(lx *lexer) stateFn {
 		lx.backup()
 		if lx.isBool() {
 			lx.emit(itemBool)
+		} else if lx.isVariable() {
+			lx.emit(itemVariable)
 		} else {
 			lx.emit(itemString)
 		}
@@ -918,6 +933,8 @@ func (itype itemType) String() string {
 		return "MapEnd"
 	case itemCommentStart:
 		return "CommentStart"
+	case itemVariable:
+		return "Variable"
 	}
 	panic(fmt.Sprintf("BUG: Unknown type '%s'.", itype.String()))
 }

diff --git a/conf/lex_test.go b/conf/lex_test.go
@@ -20,6 +20,15 @@ func expect(t *testing.T, lx *lexer, items []item) {
 	}
 }
 
+func TestPlainValue(t *testing.T) {
+	expectedItems := []item{
+		{itemKey, "foo", 1},
+		{itemEOF, "", 1},
+	}
+	lx := lex("foo")
+	expect(t, lx, expectedItems)
+}
+
 func TestSimpleKeyStringValues(t *testing.T) {
 	expectedItems := []item{
 		{itemKey, "foo", 1},
@@ -192,6 +201,20 @@ func TestDateValues(t *testing.T) {
 	expect(t, lx, expectedItems)
 }
 
+func TestVariableValues(t *testing.T) {
+	expectedItems := []item{
+		{itemKey, "foo", 1},
+		{itemVariable, "bar", 1},
+		{itemEOF, "", 1},
+	}
+	lx := lex("foo = $bar")
+	expect(t, lx, expectedItems)
+	lx = lex("foo =$bar")
+	expect(t, lx, expectedItems)
+	lx = lex("foo $bar")
+	expect(t, lx, expectedItems)
+}
+
 func TestArrays(t *testing.T) {
 	expectedItems := []item{
 		{itemKey, "foo", 1},

diff --git a/conf/parse.go b/conf/parse.go
@@ -16,6 +16,7 @@ package conf
 
 import (
 	"fmt"
+	"os"
 	"strconv"
 	"time"
 )
@@ -157,11 +158,50 @@ func (p *parser) processItem(it item) error {
 		array := p.ctx
 		p.popContext()
 		p.setValue(array)
+	case itemVariable:
+		if value, ok := p.lookupVariable(it.val); ok {
+			p.setValue(value)
+		} else {
+			return fmt.Errorf("Variable reference for '%s' on line %d can not be found.",
+				it.val, it.line)
+		}
 	}
 
 	return nil
 }
 
+// Used to map an environment value into a temporary map to pass to secondary Parse call.
+const pkey = "pk"
+
+// lookupVariable will lookup a variable reference. It will use block scoping on keys
+// it has seen before, with the top level scoping being the environment variables. We
+// ignore array contexts and only process the map contexts..
+//
+// Returns true for ok if it finds something, similar to map.
+func (p *parser) lookupVariable(varReference string) (interface{}, bool) {
+	// Loop through contexts currently on the stack.
+	for i := len(p.ctxs) - 1; i >= 0; i -= 1 {
+		ctx := p.ctxs[i]
+		// Process if it is a map context
+		if m, ok := ctx.(map[string]interface{}); ok {
+			if v, ok := m[varReference]; ok {
+				return v, ok
+			}
+		}
+	}
+
+	// If we are here, we have exhausted our context maps and still not found anything.
+	// Parse from the environment.
+	if vStr, ok := os.LookupEnv(varReference); ok {
+		// Everything we get here will be a string value, so we need to process as a parser would.
+		if vmap, err := Parse(fmt.Sprintf("%s=%s", pkey, vStr)); err == nil {
+			v, ok := vmap[pkey]
+			return v, ok
+		}
+	}
+	return nil, false
+}
+
 func (p *parser) setValue(val interface{}) {
 	// Test to see if we are on an array or a map
 

diff --git a/conf/parse_test.go b/conf/parse_test.go
@@ -1,7 +1,10 @@
 package conf
 
 import (
+	"fmt"
+	"os"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 )
@@ -32,6 +35,60 @@ func TestSimpleTopLevel(t *testing.T) {
 	test(t, "foo='1'; bar=2.2; baz=true; boo=22", ex)
 }
 
+var varSample = `
+  index = 22
+  foo = $index
+`
+
+func TestSimpleVariable(t *testing.T) {
+	ex := map[string]interface{}{
+		"index": int64(22),
+		"foo":   int64(22),
+	}
+	test(t, varSample, ex)
+}
+
+var varNestedSample = `
+  index = 22
+  nest {
+    index = 11
+    foo = $index
+  }
+  bar = $index
+`
+
+func TestNestedVariable(t *testing.T) {
+	ex := map[string]interface{}{
+		"index": int64(22),
+		"nest": map[string]interface{}{
+			"index": int64(11),
+			"foo":   int64(11),
+		},
+		"bar": int64(22),
+	}
+	test(t, varNestedSample, ex)
+}
+
+func TestMissingVariable(t *testing.T) {
+	_, err := Parse("foo=$index")
+	if err == nil {
+		t.Fatalf("Expected an error for a missing variable, got none")
+	}
+	if !strings.HasPrefix(err.Error(), "Variable reference") {
+		t.Fatalf("Wanted a variable reference err, got %q\n", err)
+	}
+}
+
+func TestEnvVariable(t *testing.T) {
+	ex := map[string]interface{}{
+		"foo": int64(22),
+	}
+	evar := "__UNIQ22__"
+	os.Setenv(evar, "22")
+	defer os.Unsetenv(evar)
+	test(t, fmt.Sprintf("foo = $%s", evar), ex)
+}
+
 var sample1 = `
 foo  {
   host {