feat: env vars in templates

resolves #743

src/console_title.go

@@ -21,8 +21,6 @@ const (
 	FullPath ConsoleTitleStyle = "path"
 	// Template allows a more powerful custom string
 	Template ConsoleTitleStyle = "template"
-
-	templateEnvRegex = `\.Env\.(?P<ENV>[^ \.}]*)`
 )
 
 func (t *consoleTitle) getConsoleTitle() string {
@@ -54,17 +52,10 @@ func (t *consoleTitle) getTemplateText() string {
 		context["Host"] = host
 	}
 
-	// load environment variables into the map
-	envVars := map[string]string{}
-	matches := findAllNamedRegexMatch(templateEnvRegex, t.config.ConsoleTitleTemplate)
-	for _, match := range matches {
-		envVars[match["ENV"]] = t.env.getenv(match["ENV"])
-	}
-	context["Env"] = envVars
-
 	template := &textTemplate{
 		Template: t.config.ConsoleTitleTemplate,
 		Context:  context,
+		Env:      t.env,
 	}
 	text, err := template.render()
 	if err != nil {

src/main.go

@@ -276,6 +276,7 @@ func getConsoleBackgroundColor(env environmentInfo, backgroundColorTemplate stri
 	template := &textTemplate{
 		Template: backgroundColorTemplate,
 		Context:  context,
+		Env:      env,
 	}
 	text, err := template.render()
 	if err != nil {

src/segment.go

@@ -23,6 +23,7 @@ type Segment struct {
 	writer              SegmentWriter
 	stringValue         string
 	active              bool
+	env                 environmentInfo
 }
 
 // SegmentTiming holds the timing context for a segment
@@ -182,6 +183,7 @@ func (segment *Segment) getColor(templates []string, defaultColor string) string
 	}
 	txtTemplate := &textTemplate{
 		Context: segment.writer,
+		Env:     segment.env,
 	}
 	for _, template := range templates {
 		txtTemplate.Template = template
@@ -211,6 +213,7 @@ func (segment *Segment) background() string {
 }
 
 func (segment *Segment) mapSegmentWithWriter(env environmentInfo) error {
+	segment.env = env
 	functions := map[SegmentType]SegmentWriter{
 		Session:       &session{},
 		Path:          &path{},

src/segment_aws.go

@@ -83,6 +83,7 @@ func (a *aws) string() string {
 	template := &textTemplate{
 		Template: segmentTemplate,
 		Context:  a,
+		Env:      a.env,
 	}
 	text, err := template.render()
 	if err != nil {

src/segment_battery.go

@@ -129,6 +129,7 @@ func (b *batt) string() string {
 	template := &textTemplate{
 		Template: segmentTemplate,
 		Context:  b,
+		Env:      b.env,
 	}
 	text, err := template.render()
 	if err != nil {

src/segment_kubectl.go

@@ -16,6 +16,7 @@ func (k *kubectl) string() string {
 	template := &textTemplate{
 		Template: segmentTemplate,
 		Context:  k,
+		Env:      k.env,
 	}
 	text, err := template.render()
 	if err != nil {

src/segment_session.go

@@ -46,6 +46,7 @@ func (s *session) enabled() bool {
 		template := &textTemplate{
 			Template: segmentTemplate,
 			Context:  s,
+			Env:      s.env,
 		}
 		var err error
 		s.templateText, err = template.render()

src/segment_time.go

@@ -24,6 +24,7 @@ func (t *tempus) enabled() bool {
 		template := &textTemplate{
 			Template: segmentTemplate,
 			Context:  t,
+			Env:      t.env,
 		}
 		var err error
 		t.templateText, err = template.render()

src/template.go

@@ -3,6 +3,8 @@ package main
 import (
 	"bytes"
 	"errors"
+	"reflect"
+	"strings"
 	"text/template"
 
 	"github.com/Masterminds/sprig"
@@ -12,23 +14,81 @@ const (
 	// Errors to show when the template handling fails
 	invalidTemplate   = "invalid template text"
 	incorrectTemplate = "unable to create text based on template"
+
+	templateEnvRegex = `\.Env\.(?P<ENV>[^ \.}]*)`
 )
 
 type textTemplate struct {
 	Template string
 	Context  interface{}
+	Env      environmentInfo
 }
 
 func (t *textTemplate) render() (string, error) {
 	tmpl, err := template.New("title").Funcs(sprig.TxtFuncMap()).Parse(t.Template)
 	if err != nil {
 		return "", errors.New(invalidTemplate)
 	}
+	if strings.Contains(t.Template, ".Env") {
+		t.loadEnvVars()
+	}
 	buffer := new(bytes.Buffer)
 	defer buffer.Reset()
 	err = tmpl.Execute(buffer, t.Context)
 	if err != nil {
 		return "", errors.New(incorrectTemplate)
 	}
-	return buffer.String(), nil
+	text := buffer.String()
+	// issue with missingkey=zero ignored for map[string]interface{}
+	// https://github.com/golang/go/issues/24963
+	text = strings.ReplaceAll(text, "<no value>", "")
+	return text, nil
+}
+
+func (t *textTemplate) loadEnvVars() {
+	context := make(map[string]interface{})
+	switch v := t.Context.(type) {
+	case map[string]interface{}:
+		context = v
+	default:
+		// we currently only support structs
+		if !t.isStruct() {
+			break
+		}
+		context = t.structToMap()
+	}
+	envVars := map[string]string{}
+	matches := findAllNamedRegexMatch(templateEnvRegex, t.Template)
+	for _, match := range matches {
+		envVars[match["ENV"]] = t.Env.getenv(match["ENV"])
+	}
+	context["Env"] = envVars
+	t.Context = context
+}
+
+func (t *textTemplate) isStruct() bool {
+	v := reflect.TypeOf(t.Context)
+	if v == nil {
+		return false
+	}
+	if v.Kind() == reflect.Ptr {
+		v = v.Elem()
+	}
+	if v.Kind() == reflect.Invalid {
+		return false
+	}
+	return v.Kind() == reflect.Struct
+}
+
+func (t *textTemplate) structToMap() map[string]interface{} {
+	context := make(map[string]interface{})
+	v := reflect.ValueOf(t.Context)
+	strct := v.Type()
+	for i := 0; i < strct.NumField(); i++ {
+		sf := strct.Field(i)
+		name := sf.Name
+		value := v.Field(i).Interface()
+		context[name] = value
+	}
+	return context
 }

src/template_test.go

@@ -68,10 +68,65 @@ func TestRenderTemplate(t *testing.T) {
 		},
 	}
 
+	env := &MockedEnvironment{}
 	for _, tc := range cases {
 		template := &textTemplate{
 			Template: tc.Template,
 			Context:  tc.Context,
+			Env:      env,
+		}
+		text, err := template.render()
+		if tc.ShouldError {
+			assert.Error(t, err)
+			continue
+		}
+		assert.Equal(t, tc.Expected, text, tc.Case)
+	}
+}
+
+func TestRenderTemplateEnvVar(t *testing.T) {
+	cases := []struct {
+		Case        string
+		Expected    string
+		Template    string
+		ShouldError bool
+		Env         map[string]string
+		Context     interface{}
+	}{
+		{
+			Case:     "map with env var",
+			Expected: "hello world",
+			Template: "{{.Env.HELLO}} {{.World}}",
+			Context:  map[string]interface{}{"World": "world"},
+			Env:      map[string]string{"HELLO": "hello"},
+		},
+		{
+			Case:     "nil struct with env var",
+			Expected: "hello world",
+			Template: "{{.Env.HELLO }} world{{ .Text}}",
+			Context:  nil,
+			Env:      map[string]string{"HELLO": "hello"},
+		},
+		{
+			Case:     "struct with env var",
+			Expected: "hello world posh",
+			Template: "{{.Env.HELLO}} world {{ .Text }}",
+			Context:  struct{ Text string }{Text: "posh"},
+			Env:      map[string]string{"HELLO": "hello"},
+		},
+		{Case: "no env var", Expected: "hello world", Template: "{{.Text}} world", Context: struct{ Text string }{Text: "hello"}},
+		{Case: "map", Expected: "hello world", Template: "{{.Text}} world", Context: map[string]interface{}{"Text": "hello"}},
+		{Case: "empty map", Expected: " world", Template: "{{.Text}} world", Context: map[string]string{}},
+	}
+	for _, tc := range cases {
+		env := &MockedEnvironment{}
+		for name, value := range tc.Env {
+			env.On("getenv", name).Return(value)
+		}
+		template := &textTemplate{
+			Template: tc.Template,
+			Context:  tc.Context,
+			Env:      env,
 		}
 		text, err := template.render()
 		if tc.ShouldError {