Merge pull request #1 from cvstom/feature/multiple-prefixes

Trim multiple trimprefixes
cvstom · Jun 15, 2020 · e64b95e · e64b95e
2 parents b0b315d + 5be06e1
commit e64b95e
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git a/golden_test.go b/golden_test.go
@@ -1226,32 +1226,32 @@ func (i Prime) IsAPrime() bool {
 
 func TestGolden(t *testing.T) {
 	for _, test := range golden {
-		runGoldenTest(t, test, false, false, false, false, false, "")
+		runGoldenTest(t, test, false, false, false, false, false, nil)
 	}
 	for _, test := range goldenJSON {
-		runGoldenTest(t, test, true, false, false, false, false, "")
+		runGoldenTest(t, test, true, false, false, false, false, nil)
 	}
 	for _, test := range goldenText {
-		runGoldenTest(t, test, false, false, false, false, true, "")
+		runGoldenTest(t, test, false, false, false, false, true, nil)
 	}
 	for _, test := range goldenYAML {
-		runGoldenTest(t, test, false, true, false, false, false, "")
+		runGoldenTest(t, test, false, true, false, false, false, nil)
 	}
 	for _, test := range goldenSQL {
-		runGoldenTest(t, test, false, false, true, false, false, "")
+		runGoldenTest(t, test, false, false, true, false, false, nil)
 	}
 	for _, test := range goldenGQL {
-		runGoldenTest(t, test, false, false, false, true, false, "")
+		runGoldenTest(t, test, false, false, false, true, false, nil)
 	}
 	for _, test := range goldenJSONAndSQL {
-		runGoldenTest(t, test, true, false, true, false, false, "")
+		runGoldenTest(t, test, true, false, true, false, false, nil)
 	}
 	for _, test := range goldenPrefix {
-		runGoldenTest(t, test, false, false, false, false, false, "Day")
+		runGoldenTest(t, test, false, false, false, false, false, []string{"Day"})
 	}
 }
 
-func runGoldenTest(t *testing.T, test Golden, generateJSON, generateYAML, generateSQL, generateGQL, generateText bool, prefix string) {
+func runGoldenTest(t *testing.T, test Golden, generateJSON, generateYAML, generateSQL, generateGQL, generateText bool, prefix []string) {
 	var g Generator
 	input := "package test\n" + test.input
 	file := test.name + ".go"

diff --git a/stringer.go b/stringer.go
@@ -52,13 +52,15 @@ var (
 	text            = flag.Bool("text", false, "if true, text marshaling methods will be generated. Default: false")
 	output          = flag.String("output", "", "output file name; default srcdir/<type>_string.go")
 	transformMethod = flag.String("transform", "noop", "enum item name transformation method. Default: noop")
-	trimPrefix      = flag.String("trimprefix", "", "transform each item name by removing a prefix. Default: \"\"")
 	lineComment     = flag.Bool("linecomment", false, "use line comment text as printed text when present")
 )
 
-var comments arrayFlags
+var (
+	comments, trimprefix arrayFlags
+)
 
 func init() {
+	flag.Var(&trimprefix, "trimprefix", "transform each item name by removing a prefix. Default: \"\"")
 	flag.Var(&comments, "comment", "comments to include in generated code, can repeat. Default: \"\"")
 }
 
@@ -127,7 +129,7 @@ func main() {
 
 	// Run generate for each type.
 	for _, typeName := range types {
-		g.generate(typeName, *json, *yaml, *sql, *gql, *text, *transformMethod, *trimPrefix, *lineComment)
+		g.generate(typeName, *json, *yaml, *sql, *gql, *text, *transformMethod, trimprefix, *lineComment)
 	}
 
 	// Format the output.
@@ -332,9 +334,11 @@ func (g *Generator) transformValueNames(values []Value, transformMethod string)
 }
 
 // trimValueNames removes a prefix from each name
-func (g *Generator) trimValueNames(values []Value, prefix string) {
+func (g *Generator) trimValueNames(values []Value, prefix []string) {
 	for i := range values {
-		values[i].name = strings.TrimPrefix(values[i].name, prefix)
+		for _, p := range prefix {
+			values[i].name = strings.TrimPrefix(values[i].name, p)
+		}
 	}
 }
 
@@ -347,7 +351,7 @@ func (g *Generator) replaceValuesWithLineComment(values []Value) {
 }
 
 // generate produces the String method for the named type.
-func (g *Generator) generate(typeName string, includeJSON, includeYAML, includeSQL, includeGQL, includeText bool, transformMethod string, trimPrefix string, lineComment bool) {
+func (g *Generator) generate(typeName string, includeJSON, includeYAML, includeSQL, includeGQL, includeText bool, transformMethod string, trimPrefix []string, lineComment bool) {
 	values := make([]Value, 0, 100)
 	for _, file := range g.pkg.files {
 		// Set the state for this run of the walker.