diff --git a/config.go b/config.go index 97d293b..d4254e4 100644 --- a/config.go +++ b/config.go @@ -10,11 +10,12 @@ import ( var conf = &Config{} type Config struct { - Encoding EncoderName - TypeID TypeIDName - Debug bool - DstDir string - ModPath string + Encoding EncoderName + TypeID TypeIDName + Debug bool + DstDir string + ModPath string + RemoveAccountSuffix bool } func GetConfig() *Config { diff --git a/generator_test.go b/generator_test.go index 4602f9d..2b64c39 100644 --- a/generator_test.go +++ b/generator_test.go @@ -7,6 +7,7 @@ import ( . "github.com/dave/jennifer/jen" "github.com/davecgh/go-spew/spew" "github.com/stretchr/testify/require" + "github.com/stretchr/testify/assert" ) func Test_genTypeName(t *testing.T) { @@ -271,3 +272,23 @@ func Test_IdlAccountItemSlice_Walk(t *testing.T) { require.Equal(t, expectedAccountNames, gotAccountNames) require.Equal(t, expectedIndexes, gotIndexes) } + +func TestFormatAccountAccessorName(t *testing.T) { + t.Run("default config", func(t *testing.T) { + assert.Equal(t, "GetFooAccount", formatAccountAccessorName("Get", "Foo")) + assert.Equal(t, "GetFooAccountAccount", formatAccountAccessorName("Get", "FooAccount")) + }) + + t.Run("remove config on", func(t *testing.T) { + oldConf := GetConfig() + defer func() { + conf = oldConf + }() + conf = &Config{ + RemoveAccountSuffix: true, + } + + assert.Equal(t, "GetFooAccount", formatAccountAccessorName("Get", "Foo")) + assert.Equal(t, "GetFooAccount", formatAccountAccessorName("Get", "FooAccount")) + }) +} diff --git a/main.go b/main.go index 7bcccfe..029599e 100644 --- a/main.go +++ b/main.go @@ -28,6 +28,7 @@ func main() { flag.Var(&filenames, "src", "Path to source; can use multiple times.") flag.StringVar(&conf.DstDir, "dst", generatedDir, "Destination folder") flag.BoolVar(&conf.Debug, "debug", false, "debug mode") + flag.BoolVar(&conf.RemoveAccountSuffix, "remove-account-suffix", false, "Remove \"Account\" suffix from accessors (if leads to duplication, e.g. \"SetFooAccountAccount\")") flag.StringVar((*string)(&conf.Encoding), "codec", string(EncodingBorsh), "Choose codec") flag.StringVar((*string)(&conf.TypeID), "type-id", string(TypeIDAnchor), "Choose typeID kind") @@ -1634,16 +1635,16 @@ func genProgramBoilerplate(idl IDL) (*File, error) { // formatAccountAccessorName formats a name for a function that // either gets or sets an account. -// If the name already has a "Account" suffix, then another "Account" suffix -// is NOT added. +// If the RemoveAccountSuffix config flag is set, and the name already +// has an "Account" suffix, then another "Account" suffix is NOT added. // E.g. ("Set", "Foo") => "SetFooAccount" // E.g. ("Set", "BarAccount") => "SetBarAccount" func formatAccountAccessorName(prefix, name string) string { endsWithAccount := strings.HasSuffix(strings.ToLower(name), "account") - if endsWithAccount { - return prefix + name + if !conf.RemoveAccountSuffix || !endsWithAccount { + return prefix + name + "Account" } - return prefix + name + "Account" + return prefix + name } func treeFindLongestNameFromFields(fields []IdlField) (ln int) {