Skip to content

Commit 8b37e45

Browse files
committed
Go: Add helpers for defined enum types
Adds basic string conversion / lookup helpers for compiled enums Signed-off-by: Kendall Goto <kendall@level.co>
1 parent 3603e70 commit 8b37e45

File tree

1 file changed

+47
-3
lines changed

1 file changed

+47
-3
lines changed

shared/src/main/scala/io/kaitai/struct/languages/GoCompiler.scala

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -536,15 +536,59 @@ class GoCompiler(typeProvider: ClassTypeProvider, config: RuntimeConfig)
536536

537537
out.dec
538538
out.puts(")")
539+
539540
// Inspired by https://gist.github.com/bgadrian/cb8b9344d9c66571ef331a14eb7a2e80
540-
val mapEntriesStr = enumColl.map { case (id, _) => s"$id: {}" }.mkString(", ")
541-
out.puts(s"var values_$fullEnumNameStr = map[$fullEnumNameStr]struct{}{$mapEntriesStr}")
541+
val mapEntriesStr = enumColl.map { case (id, label) => s"$id: \"${type2class(label.name)}\"" }.mkString(", ")
542+
out.puts(s"var valueNames_$fullEnumNameStr = map[$fullEnumNameStr]string{$mapEntriesStr}")
543+
542544
out.puts(s"func (v $fullEnumNameStr) isDefined() bool {")
543545
out.inc
544-
out.puts(s"_, ok := values_$fullEnumNameStr[v]")
546+
out.puts(s"_, ok := valueNames_$fullEnumNameStr[v]")
545547
out.puts("return ok")
546548
out.dec
547549
out.puts("}")
550+
551+
out.puts(s"func (v $fullEnumNameStr) IsA$fullEnumNameStr() bool {")
552+
out.inc
553+
out.puts(s"return v.isDefined()")
554+
out.dec
555+
out.puts("}")
556+
557+
out.puts(s"func (v $fullEnumNameStr) String() string {")
558+
out.inc
559+
out.puts(s"name, ok := valueNames_$fullEnumNameStr[v]")
560+
out.puts("if ok {")
561+
out.inc
562+
out.puts("return name")
563+
out.dec
564+
out.puts("}")
565+
out.puts(s"return \"${fullEnumNameStr}(\" + string(v) + \")\"")
566+
out.dec
567+
out.puts("}")
568+
569+
out.puts(s"func ${fullEnumNameStr}Strings() []string {")
570+
out.inc
571+
out.puts(s"strings := make([]string, 0, len(valueNames_$fullEnumNameStr))")
572+
out.puts(s"for _, name := range valueNames_$fullEnumNameStr {")
573+
out.inc
574+
out.puts("strings = append(strings, name)")
575+
out.dec
576+
out.puts("}")
577+
out.puts("return strings")
578+
out.dec
579+
out.puts("}")
580+
581+
out.puts(s"func ${fullEnumNameStr}Values() []${fullEnumNameStr} {")
582+
out.inc
583+
out.puts(s"values := make([]${fullEnumNameStr}, 0, len(valueNames_$fullEnumNameStr))")
584+
out.puts(s"for value := range valueNames_$fullEnumNameStr {")
585+
out.inc
586+
out.puts("values = append(values, value)")
587+
out.dec
588+
out.puts("}")
589+
out.puts("return values")
590+
out.dec
591+
out.puts("}")
548592
}
549593

550594
override def classToString(toStringExpr: Ast.expr): Unit = {

0 commit comments

Comments
 (0)