From a1a049db8d267519553ec0ac2536e83b75a241a9 Mon Sep 17 00:00:00 2001
From: maskpp <maskpp266@gmail.com>
Date: Sun, 24 Sep 2023 21:14:24 +0800
Subject: [PATCH] Allow abigen to use custom templates

---
 accounts/abi/bind/template.go |  7 +++++++
 cmd/abigen/main.go            | 21 +++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/accounts/abi/bind/template.go b/accounts/abi/bind/template.go
index 95dc13cc18..4b6f82f1ae 100644
--- a/accounts/abi/bind/template.go
+++ b/accounts/abi/bind/template.go
@@ -78,6 +78,13 @@ var tmplSource = map[Lang]string{
 	LangGo: tmplSourceGo,
 }
 
+// SetTmplSource supports this func in order to set special template file.
+func SetTmplSource(lang Lang, source string) {
+	if _, ok := tmplSource[lang]; ok {
+		tmplSource[lang] = source
+	}
+}
+
 // tmplSourceGo is the Go source template that the generated Go contract binding
 // is based on.
 const tmplSourceGo = `
diff --git a/cmd/abigen/main.go b/cmd/abigen/main.go
index 221f45c078..3b2094b093 100644
--- a/cmd/abigen/main.go
+++ b/cmd/abigen/main.go
@@ -72,6 +72,14 @@ var (
 		Name:  "alias",
 		Usage: "Comma separated aliases for function and event renaming, e.g. original1=alias1, original2=alias2",
 	}
+	contractFlag = cli.StringFlag{
+		Name:  "contract",
+		Usage: "Name of the contract to generate the bindings for",
+	}
+	tmplFlag = cli.StringFlag{
+		Name:  "tmpl",
+		Usage: "Template file if a user wants to customize",
+	}
 )
 
 var app = flags.NewApp("Ethereum ABI wrapper code generator")
@@ -182,6 +190,10 @@ func abigen(c *cli.Context) error {
 			// fully qualified name is of the form <solFilePath>:<type>
 			nameParts := strings.Split(name, ":")
 			typeName := nameParts[len(nameParts)-1]
+			// If a contract name is provided then ignore all other contracts
+			if c.IsSet(contractFlag.Name) && c.String(contractFlag.Name) != typeName {
+				continue
+			}
 			if exclude != nil && exclude.Matches(name) {
 				fmt.Fprintf(os.Stderr, "excluding: %v\n", name)
 				continue
@@ -215,6 +227,15 @@ func abigen(c *cli.Context) error {
 			aliases[match[1]] = match[2]
 		}
 	}
+	// Set customize template file.
+	if c.IsSet(tmplFlag.Name) {
+		tmplFile := c.String(tmplFlag.Name)
+		data, err := os.ReadFile(tmplFile)
+		if err != nil {
+			utils.Fatalf("Failed to read template file: %v", err)
+		}
+		bind.SetTmplSource(lang, string(data))
+	}
 	// Generate the contract binding
 	code, err := bind.Bind(types, abis, bins, sigs, c.String(pkgFlag.Name), lang, libs, aliases)
 	if err != nil {