expand package for Go to handle ${VAR:-foo} style expansions

expand.go

// Package expand is used to handle basic ${VAR:-DEFAULT} style text interpolations.
//
// It currently only supports ASCII variable names, since you probably won't see variations on that
// in the wild.

package expand

import (
	"bytes"
	"strings"
	"unicode"
)

// LookupFunc is any function that looks up a variable with the given name and returns its value and
// whether it's defined.
type LookupFunc func(key string) (value string, ok bool)

// GetFunc converts a function that returns a value for a string and returns a LookupFunc for it.
// The LookupFunc returns that a value is not defined if fn returns an empty string.
func GetFunc(fn func(string) string) LookupFunc {
	return func(key string) (string, bool) {
		v := fn(key)
		return v, v != ""
	}
}

// Expand expands an input string using shell-like expansions $VAR, ${VAR}, ${VAR:-if_undefined}, and
// ${VAR:+if_defined} and returns the result of the expansions.
//
// Cases of "$VAR" and "${VAR}" are replaced with empty strings if VAR is not defined.
// "${VAR:-if_undefined}" replaces VAR with the text "if_undefined" if VAR is not defined.
// "${VAR:+if_defined}" replaces VAR with the text "if_defined" if VAR is defined.
func Expand(in string, fn LookupFunc) string {
	return ExpandCustom(in, '$', false, fn)
}

// ExpandCustom expands an input string using shell-like expansions <lead>VAR, <lead>VAR},
// <lead>VAR:-if_undefined}, and <lead>VAR:+if_defined} and returns the result of the expansions.
// If curlyOnly is true, <lead>VAR expansions are ignored.
//
// Cases of "<lead>VAR" and "<lead>{VAR}" are replaced with empty strings if VAR is not defined.
// "<lead>{VAR:-if_undefined}" replaces VAR with the text "if_undefined" if VAR is not defined.
// "<lead>{VAR:+if_defined}" replaces VAR with the text "if_defined" if VAR is defined.
func ExpandCustom(in string, lead byte, curlyOnly bool, fn LookupFunc) string {
	var buf bytes.Buffer
	buf.Grow(len(in) * 2)
	_, ex := parse(in, -1, idChar(lead), curlyOnly)
	ex.expand(&buf, fn)
	return buf.String()
}

type expansion struct {
	value    string
	kind     exKind
	children []*expansion
}

type exKind int

const (
	exPrefix   exKind = iota // Literal text
	exVariable               // $VAR ${VAR} ${VAR:-UNDEFINED TEXT}
	exDefined                // ${VAR:+DEFINED TEXT}
)

func (e *expansion) expand(w *bytes.Buffer, fn LookupFunc) {
	if e == nil {
		return
	}

	switch e.kind {
	case exPrefix:
		if v := e.value; v != "" {
			w.WriteString(v)
		}
	case exVariable:
		if v, ok := fn(e.value); ok {
			if v != "" {
				w.WriteString(v)
			}
			return
		}
	case exDefined:
		if _, ok := fn(e.value); !ok {
			return
		}
	}

	for _, e := range e.children {
		e.expand(w, fn)
	}
}

type idChar byte

func (r idChar) is(c byte) bool {
	return !(c == ':' || c == '}' || c == byte(r) || unicode.IsSpace(rune(c)))
}

func parse(in string, until rune, lead idChar, curlyOnly bool) (n int, root *expansion) {
	var (
		i       int
		c       byte
		consume = func(ex *expansion) {
			switch {
			case len(ex.value) == 0 && ex.kind == exPrefix:
			case root == nil && ex.kind == exPrefix:
				root = ex
			case root == nil:
				root = &expansion{children: []*expansion{ex}}
			case ex.kind == exPrefix && len(root.children) == 0:
				root.value += ex.value
			default:
				root.children = append(root.children, ex)
			}
			n += len(in[:i])
			in = in[i:]
			i = 0
		}
	)

reset:
	for ; i < len(in); i++ {
		c = in[i]
		switch {
		case until != -1 && rune(c) == until:
			in = in[:i]
			goto done
		case curlyOnly && i+3 < len(in) && c == '\\' && // Treat \\LEAD{ as literal text \ followed by an expansion
			in[i+1] == '\\' && in[i+2] == byte(lead) && in[i+3] == '{':
			consume(&expansion{value: in[:i]})
			in = in[i+1:]
			continue
		case curlyOnly && i+2 < len(in) && c == '\\' && // Treat \LEAD{ as literal text LEAD{
			in[i+1] == byte(lead) && in[i+2] == '{':
			consume(&expansion{value: in[:i]})
			in = in[i+1:]
			continue
		case c != byte(lead):
			continue
		case len(in) <= i+1:
			// consume(&expansion{value: in[:i], kind: exPrefix})
		case !curlyOnly && in[i+1] == byte(lead):
			consume(&expansion{value: in[:i]})
			in = in[i+1:]
			continue
		case curlyOnly && c == byte(lead) && in[i+1] != '{':
			continue
		case i > 0:
			consume(&expansion{value: in[:i], kind: exPrefix})
		}

		i++
		break
	}

	if i >= len(in) {
		goto done
	} else if in[i] == '{' {
		i++
	} else {
		head := i
		for ; i < len(in); i++ {
			if lead.is(in[i]) {
				continue
			} else if i != head {
				consume(&expansion{value: in[head:i], kind: exVariable})
			}
			goto reset
		}

		consume(&expansion{value: in[head:i], kind: exVariable})
		goto done
	}

	for head := i; i < len(in); i++ {
		switch c = in[i]; c {
		case '}':
			tail := i
			i++
			if head == tail {
			} else {
				consume(&expansion{value: in[head:tail], kind: exVariable})
			}
			goto reset

		case ':':
			var (
				n       int
				back    *expansion
				tail    = i
				kind    = exVariable
				varname string
			)
			i++
			if head == tail { // No variable name
				goto reset
			}

			varname = strings.TrimSpace(in[head:tail])
			switch in[i] {
			case '-':
			case '+':
				kind = exDefined
			default:
				goto reset
			}

			i++
			n, back = parse(in[i:], '}', lead, curlyOnly)
			i += n
			if i < len(in) && in[i] == '}' {
				i++
				consume(&expansion{value: varname, kind: kind, children: []*expansion{back}})
			} else {
				consume(&expansion{value: in[:tail+2], kind: exPrefix})
				consume(back)
			}

			goto reset
		}

		if !lead.is(c) {
			consume(&expansion{value: in[:i]})
			goto reset
		}
	}

done:
	if len(in) > 0 {
		consume(&expansion{value: in})
	}

	if root != nil && root.value == "" && len(root.children) == 1 {
		root = root.children[0]
	}

	return n, root
}

expand_test.go

package expand

import "testing"

func TestExpansions(t *testing.T) {
	lookup := func(key string) (value string, ok bool) {
		return "text", key != "none"
	}

	type testCase struct {
		Name string
		In   string
		Want string
	}

	cases := []testCase{
		{
			"Empty",
			"",
			"",
		},
		{
			"AllDollarSigns",
			"$$$$$$$$$",
			"$$$$$",
		},
		{
			"Literal",
			"Literal",
			"Literal",
		},
		{
			"Delimiter",
			"${}{}",
			"${}{}",
		},
		{
			"EscapeDelimiterBraced",
			"$${key}",
			"${key}",
		},
		{
			"ComplexTrailingLiteral",
			"${key}literal",
			"textliteral",
		},
		{
			"ComplexLeadingLiteral",
			"literal${key}",
			"literaltext",
		},
		{
			"ComplexPaddedLiteral",
			"literal${key}literal",
			"literaltextliteral",
		},
		{
			"ComplexPaddedLiteralMulti",
			"${key:-foo}literal${key}literal${key:+bar}",
			"textliteraltextliteralbar",
		},
		{
			"EscapeDelimiterBracedUndefined",
			"$${none:-undefined}${none:-undefined}",
			"${none:-undefined}undefined",
		},
		{
			"EscapeDelimiterBracedDefined",
			"$${key:+defined}${key:+defined}",
			"${key:+defined}defined",
		},
		{
			"NestedInterpolations",
			"$$${key:+defined: ${key} ${none:-undefined}}$$",
			"$defined: text undefined$",
		},
		{
			"NestedInterpolations",
			"$$${key:-undefined: ${key} ${none:-undefined}}$$",
			"$text$",
		},
		{
			"ExpansionTypes",
			"literal $key $none ${key} ${none} ${key:-undefined} ${none:-undefined} ${key:+defined} ${none:+defined}",
			"literal text  text  text undefined defined ",
		},
		{
			"IncompleteExpansionNone",
			"${key}${key:${none:-alice}",
			"text${key:alice",
		},
		{
			"IncompleteExpansionUndefined",
			"${key}${key:-${none:-bob}",
			"text${key:-bob",
		},
		{
			"IncompleteExpansionDefined",
			"${key}${key:+${none:-carol}",
			"text${key:+carol",
		},
		{
			"OnlyVariable",
			"$variable",
			"text",
		},
		{
			"NoVariable",
			"${}${ }${:}${:+}${:-}$ ",
			"${}${ }${:}${:+}${:-}$ ",
		},
		{
			"AllPrefixes",
			"Foo${ }Bar",
			"Foo${ }Bar",
		},
		{
			"UnusualCharacters",
			"${!@#%^&*()[]{.,/<>?;'\"|-+=~`é}}",
			"text}",
		},
	}

	for _, c := range cases {
		c := c
		t.Run(c.Name, func(t *testing.T) {
			if got := Expand(c.In, lookup); got != c.Want {
				t.Fatalf("Expand(%q) = %q; want %q", c.In, got, c.Want)
			}
		})
	}
}

func TestCurlyOnlyExpansions(t *testing.T) {
	lookup := func(key string) (value string, ok bool) {
		return "text", key != "none"
	}

	type testCase struct {
		Name string
		In   string
		Want string
	}

	cases := []testCase{
		{
			"Empty",
			"",
			"",
		},
		{
			"AllHashSigns",
			"#########",
			"#########",
		},
		{
			"Literal",
			"Literal",
			"Literal",
		},
		{
			"Delimiter",
			"#{}{}",
			"#{}{}",
		},
		{
			"EscapeDelimiterBraced",
			"\\#{key}",
			"#{key}",
		},
		{
			"ComplexTrailingLiteral",
			"#{key}literal",
			"textliteral",
		},
		{
			"ComplexLeadingLiteral",
			"literal#{key}",
			"literaltext",
		},
		{
			"ComplexPaddedLiteral",
			"literal#{key}literal",
			"literaltextliteral",
		},
		{
			"ComplexPaddedLiteralMulti",
			"#{key:-foo}literal#{key}literal#{key:+bar}",
			"textliteraltextliteralbar",
		},
		{
			"EscapeDelimiterBracedUndefined",
			"\\#{none:-undefined}#{none:-undefined}",
			"#{none:-undefined}undefined",
		},
		{
			"EscapeDelimiterBracedDefined",
			"\\#{key:+defined}#{key:+defined}",
			"#{key:+defined}defined",
		},
		{
			"EscapedNestedInterpolations",
			"#\\#{key:+defined: #{key} #{none:-undefined}}\\\\#{key:+ defined}##",
			"##{key:+defined: text undefined}\\ defined##",
		},
		{
			"NestedInterpolations_1",
			"##{key:+defined: #{key} #{none:-undefined}}##",
			"#defined: text undefined##",
		},
		{
			"NestedInterpolations_2",
			"###{key:-undefined: #{key} #{none:-undefined}}##",
			"##text##",
		},
		{
			"ExpansionTypes",
			"literal #key #none #{key} #{none} #{key:-undefined} #{none:-undefined} #{key:+defined} #{none:+defined}",
			"literal #key #none text  text undefined defined ",
		},
		{
			"IncompleteExpansionNone",
			"#{key}#{key:#{none:-alice}",
			"text#{key:alice",
		},
		{
			"IncompleteExpansionUndefined",
			"#{key}#{key:-#{none:-bob}",
			"text#{key:-bob",
		},
		{
			"IncompleteExpansionDefined",
			"#{key}#{key:+#{none:-carol}",
			"text#{key:+carol",
		},
		{
			"OnlyVariable",
			"#variable",
			"#variable",
		},
		{
			"NoVariable",
			"#{}#{ }#{:}#{:+}#{:-}# ",
			"#{}#{ }#{:}#{:+}#{:-}# ",
		},
		{
			"AllPrefixes",
			"Foo#{ }Bar",
			"Foo#{ }Bar",
		},
		{
			"UnusualCharacters",
			"#{!@$%^&*()[]{.,/<>?;'\"|-+=~`é}}",
			"text}",
		},
	}

	for _, c := range cases {
		c := c
		t.Run(c.Name, func(t *testing.T) {
			if got := ExpandCustom(c.In, '#', true, lookup); got != c.Want {
				t.Fatalf("Expand(%q, '#', true) = %q; want %q", c.In, got, c.Want)
			}
		})
	}
}

func TestUnicodeExpand(t *testing.T) {
	const want = "你好"
	const wantKey = "မင်္ဂလာပါ"
	const input = "${" + wantKey + "}"
	gotKey := false

	fn := func(key string) (string, bool) {
		if gotKey = wantKey == key; gotKey {
			return want, true
		}
		return "", false
	}

	got := Expand(input, fn)
	if got != want {
		t.Fatalf("Expand(%q) = %q; want %q", input, got, want)
	}
}

func TestGetFunc(t *testing.T) {
	want, wantOK := "value", true
	got, gotOK := GetFunc(func(string) string { return "value" })("key")
	if got != want || gotOK != wantOK {
		t.Errorf("GetFunc(-> value) = %q, %t; want %q, %t", got, gotOK, want, wantOK)
	}

	want, wantOK = "", false
	got, gotOK = GetFunc(func(string) string { return "" })("key")
	if got != want || gotOK != wantOK {
		t.Errorf("GetFunc(-> value) = %q, %t; want %q, %t", got, gotOK, want, wantOK)
	}
}