From a40ea91a54977a826cd8b6bd653681da1e883574 Mon Sep 17 00:00:00 2001
From: Marty Zalega <marty@zalega.me>
Date: Tue, 16 Jul 2024 20:00:12 +1000
Subject: [PATCH] Add helper functions to script templates

---
 README.md           |  8 ++++++++
 command_set.go      |  2 +-
 command_set_test.go | 24 ++++++++++++++++++++++++
 utils.go            | 23 +++++++++++++++++++++++
 4 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 53e3f60..8d21710 100644
--- a/README.md
+++ b/README.md
@@ -277,6 +277,14 @@ The expression to reference an input value. ie. '{{ .Input.my_input }}'
 
 The expression to reference an environment variable. ie. '{{ .Env.HOME }}'
 
+### input "input_name"
+
+A function to retrieve the input by its name. ie. '{{input "my_input"}}'
+
+### env "variable_name"
+
+A function to retrieve the environment variable by its name. ie. '{{env "HOME"}}'
+
 ## Example config with single command
 
 ```yaml
diff --git a/command_set.go b/command_set.go
index 55a8886..e4cd305 100644
--- a/command_set.go
+++ b/command_set.go
@@ -103,7 +103,7 @@ func (cs CommandSet) RenderScript(data TemplateData) (string, error) {
 			continue
 		}
 		if tmpl == nil {
-			tmpl = template.New(command.Name)
+			tmpl = template.New(command.Name).Funcs(data.Funcs())
 		} else {
 			tmpl = tmpl.New(command.Name)
 		}
diff --git a/command_set_test.go b/command_set_test.go
index a69ccfd..9b6d588 100644
--- a/command_set_test.go
+++ b/command_set_test.go
@@ -321,6 +321,30 @@ func TestCommandSetRenderScript(t *testing.T) {
 		assert.NoError(t, err, "CommandSet.RenderScript() returned an unexpected error")
 		assert.Equal(t, expected, actual, "CommandSet.RenderScript() returned unexpected result")
 	})
+	t.Run("helper functions", func(t *testing.T) {
+		data := TemplateData{
+			Input: map[string]any{
+				"A": "a",
+				"B": "b",
+			},
+			Env: map[string]string{
+				"C": "c",
+				"D": "d",
+			},
+		}
+		cs := CommandSet{
+			Commands: []ConfigCommand{
+				{
+					Name: "foobar",
+					Run:  "echo {{input \"A\"}} {{env \"C\"}}",
+				},
+			},
+		}
+		expected := "echo a c"
+		actual, err := cs.RenderScript(data)
+		assert.NoError(t, err, "CommandSet.RenderScript() returned an unexpected error")
+		assert.Equal(t, expected, actual, "CommandSet.RenderScript() returned unexpected result")
+	})
 }
 
 func TestCommandSetRenderScriptToTemp(t *testing.T) {
diff --git a/utils.go b/utils.go
index f509bbe..ceeb783 100644
--- a/utils.go
+++ b/utils.go
@@ -13,6 +13,29 @@ type TemplateData struct {
 	Env   map[string]string
 }
 
+func (td TemplateData) getInput(name string) any {
+	if v, ok := td.Input[name]; ok {
+		return v
+	} else {
+		return nil
+	}
+}
+
+func (td TemplateData) getEnv(name string) any {
+	if v, ok := td.Env[name]; ok {
+		return v
+	} else {
+		return nil
+	}
+}
+
+func (td *TemplateData) Funcs() template.FuncMap {
+	return template.FuncMap{
+		"input": td.getInput,
+		"env":   td.getEnv,
+	}
+}
+
 func NewTemplateData(input map[string]any, env []string) TemplateData {
 	safeInputs := make(map[string]any)
 	for name, value := range input {