From 79fa990205d3ab5049a5c028cc1a6d890e17a976 Mon Sep 17 00:00:00 2001 From: Jan De Dobbeleer Date: Wed, 26 May 2021 21:12:58 +0200 Subject: [PATCH] feat: env vars in templates resolves #743 --- src/console_title.go | 11 +------- src/main.go | 1 + src/segment.go | 3 ++ src/segment_aws.go | 1 + src/segment_battery.go | 1 + src/segment_kubectl.go | 1 + src/segment_session.go | 1 + src/segment_time.go | 1 + src/template.go | 62 +++++++++++++++++++++++++++++++++++++++++- src/template_test.go | 55 +++++++++++++++++++++++++++++++++++++ 10 files changed, 126 insertions(+), 11 deletions(-) diff --git a/src/console_title.go b/src/console_title.go index 9ae270b4..e3719ea5 100644 --- a/src/console_title.go +++ b/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[^ \.}]*)` ) 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 { diff --git a/src/main.go b/src/main.go index d3e6a109..7ae8411e 100644 --- a/src/main.go +++ b/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 { diff --git a/src/segment.go b/src/segment.go index acdece23..690de154 100644 --- a/src/segment.go +++ b/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{}, diff --git a/src/segment_aws.go b/src/segment_aws.go index aa3f3788..5864bd40 100644 --- a/src/segment_aws.go +++ b/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 { diff --git a/src/segment_battery.go b/src/segment_battery.go index 1811c9a5..28f7e3ac 100644 --- a/src/segment_battery.go +++ b/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 { diff --git a/src/segment_kubectl.go b/src/segment_kubectl.go index 3dfdf8dc..ffa543b4 100644 --- a/src/segment_kubectl.go +++ b/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 { diff --git a/src/segment_session.go b/src/segment_session.go index ed7cfca7..7a0b8e02 100644 --- a/src/segment_session.go +++ b/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() diff --git a/src/segment_time.go b/src/segment_time.go index f537c529..1404d5c0 100644 --- a/src/segment_time.go +++ b/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() diff --git a/src/template.go b/src/template.go index bbeee660..a17f3023 100644 --- a/src/template.go +++ b/src/template.go @@ -3,6 +3,8 @@ package main import ( "bytes" "errors" + "reflect" + "strings" "text/template" "github.com/Masterminds/sprig" @@ -12,11 +14,14 @@ const ( // Errors to show when the template handling fails invalidTemplate = "invalid template text" incorrectTemplate = "unable to create text based on template" + + templateEnvRegex = `\.Env\.(?P[^ \.}]*)` ) type textTemplate struct { Template string Context interface{} + Env environmentInfo } func (t *textTemplate) render() (string, error) { @@ -24,11 +29,66 @@ func (t *textTemplate) render() (string, error) { 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, "", "") + 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 } diff --git a/src/template_test.go b/src/template_test.go index dea8fa51..8e0bccf8 100644 --- a/src/template_test.go +++ b/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 {