kmsg: add TextMarshaler/TextUnmarshaler to enums

twmb · Mar 10, 2022 · 806cf53 · 806cf53
1 parent a8d2d7b
commit 806cf53
Show file tree

Hide file tree

Showing 3 changed files with 150 additions and 0 deletions.
diff --git a/generate/gen.go b/generate/gen.go
@@ -799,6 +799,22 @@ func (e Enum) WriteParseFunc(l *LineWriter) {
 	l.Write("}")
 }
 
+func (e Enum) WriteUnmarshalTextFunc(l *LineWriter) {
+	l.Write("// UnmarshalText implements encoding.TextUnmarshaler.")
+	l.Write("func (e *%s) UnmarshalText(text []byte) error {", e.Name)
+	l.Write("v, err := Parse%s(string(text))", e.Name)
+	l.Write("*e = v")
+	l.Write("return err")
+	l.Write("}")
+}
+
+func (e Enum) WriteMarshalTextFunc(l *LineWriter) {
+	l.Write("// MarshalText implements encoding.TextMarshaler.")
+	l.Write("func (e %s) MarshalText() (text []byte, err error) {", e.Name)
+	l.Write("return []byte(e.String()), nil")
+	l.Write("}")
+}
+
 func strnorm(s string) string {
 	s = strings.ReplaceAll(s, ".", "")
 	s = strings.ReplaceAll(s, "_", "")

diff --git a/generate/main.go b/generate/main.go
@@ -529,6 +529,8 @@ func main() {
 		e.WriteStringsFunc(l)
 		e.WriteParseFunc(l)
 		e.WriteConsts(l)
+		e.WriteMarshalTextFunc(l)
+		e.WriteUnmarshalTextFunc(l)
 	}
 
 	writeStrnorm(l)

diff --git a/pkg/kmsg/generated.go b/pkg/kmsg/generated.go