From 19df238b7792e8b7c7cf1889583ed45d786e1c7b Mon Sep 17 00:00:00 2001
From: Ryooooooga <eial5q265e5@gmail.com>
Date: Fri, 30 Sep 2022 21:10:56 +0900
Subject: [PATCH] feat: allow `OSCommand.Quote` to be invoked within a custom
 command

---
 docs/Custom_Command_Keybindings.md              |  2 +-
 .../services/custom_commands/handler_creator.go |  7 ++++++-
 pkg/gui/services/custom_commands/resolver.go    | 17 -----------------
 pkg/utils/template.go                           |  4 ++--
 4 files changed, 9 insertions(+), 21 deletions(-)

diff --git a/docs/Custom_Command_Keybindings.md b/docs/Custom_Command_Keybindings.md
index ca966799ec2..09708d54c25 100644
--- a/docs/Custom_Command_Keybindings.md
+++ b/docs/Custom_Command_Keybindings.md
@@ -8,7 +8,7 @@ customCommands:
     command: 'hub browse -- "commit/{{.SelectedLocalCommit.Sha}}"'
     context: 'commits'
   - key: 'a'
-    command: "git {{if .SelectedFile.HasUnstagedChanges}} add {{else}} reset {{end}} {{.SelectedFile.Name}}"
+    command: "git {{if .SelectedFile.HasUnstagedChanges}} add {{else}} reset {{end}} {{.SelectedFile.Name | Quote}}"
     context: 'files'
     description: 'toggle file staged'
   - key: 'C'
diff --git a/pkg/gui/services/custom_commands/handler_creator.go b/pkg/gui/services/custom_commands/handler_creator.go
index 839f70adfdf..9b30041e353 100644
--- a/pkg/gui/services/custom_commands/handler_creator.go
+++ b/pkg/gui/services/custom_commands/handler_creator.go
@@ -2,6 +2,7 @@ package custom_commands
 
 import (
 	"strings"
+	"text/template"
 
 	"github.com/jesseduffield/generics/slices"
 	"github.com/jesseduffield/lazygit/pkg/commands"
@@ -166,7 +167,11 @@ func (self *HandlerCreator) getResolveTemplateFn(form map[string]string, promptR
 		Form:            form,
 	}
 
-	return func(templateStr string) (string, error) { return utils.ResolveTemplate(templateStr, objects) }
+	funcs := template.FuncMap{
+		"Quote": self.os.Quote,
+	}
+
+	return func(templateStr string) (string, error) { return utils.ResolveTemplate(templateStr, objects, funcs) }
 }
 
 func (self *HandlerCreator) finalHandler(customCommand config.CustomCommand, sessionState *SessionState, promptResponses []string, form map[string]string) error {
diff --git a/pkg/gui/services/custom_commands/resolver.go b/pkg/gui/services/custom_commands/resolver.go
index 4702d36c4cf..4c32a51caac 100644
--- a/pkg/gui/services/custom_commands/resolver.go
+++ b/pkg/gui/services/custom_commands/resolver.go
@@ -1,9 +1,6 @@
 package custom_commands
 
 import (
-	"bytes"
-	"text/template"
-
 	"github.com/jesseduffield/lazygit/pkg/common"
 	"github.com/jesseduffield/lazygit/pkg/config"
 )
@@ -110,17 +107,3 @@ type CustomCommandObject struct {
 	PromptResponses []string
 	Form            map[string]string
 }
-
-func ResolveTemplate(templateStr string, object interface{}) (string, error) {
-	tmpl, err := template.New("template").Parse(templateStr)
-	if err != nil {
-		return "", err
-	}
-
-	var buf bytes.Buffer
-	if err := tmpl.Execute(&buf, object); err != nil {
-		return "", err
-	}
-
-	return buf.String(), nil
-}
diff --git a/pkg/utils/template.go b/pkg/utils/template.go
index 41388ae3a2e..d98a68b99da 100644
--- a/pkg/utils/template.go
+++ b/pkg/utils/template.go
@@ -6,8 +6,8 @@ import (
 	"text/template"
 )
 
-func ResolveTemplate(templateStr string, object interface{}) (string, error) {
-	tmpl, err := template.New("template").Option("missingkey=error").Parse(templateStr)
+func ResolveTemplate(templateStr string, object interface{}, funcs template.FuncMap) (string, error) {
+	tmpl, err := template.New("template").Funcs(funcs).Option("missingkey=error").Parse(templateStr)
 	if err != nil {
 		return "", err
 	}