Skip to content

Commit 56f208d

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 56f208d

File tree

1 file changed

+48
-3
lines changed

1 file changed

+48
-3
lines changed

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

Lines changed: 48 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -524,6 +524,7 @@ class GoCompiler(typeProvider: ClassTypeProvider, config: RuntimeConfig)
524524
override def enumDeclaration(curClass: List[String], enumName: String, enumColl: Seq[(Long, EnumValueSpec)]): Unit = {
525525
val fullEnumName: List[String] = curClass ++ List(enumName)
526526
val fullEnumNameStr = types2class(fullEnumName)
527+
importList.add("fmt")
527528

528529
out.puts
529530
out.puts(s"type $fullEnumNameStr int")
@@ -536,15 +537,59 @@ class GoCompiler(typeProvider: ClassTypeProvider, config: RuntimeConfig)
536537

537538
out.dec
538539
out.puts(")")
540+
539541
// 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}")
542+
val mapEntriesStr = enumColl.map { case (id, label) => s"$id: \"${type2class(label.name)}\"" }.mkString(", ")
543+
out.puts(s"var valueNames_$fullEnumNameStr = map[$fullEnumNameStr]string{$mapEntriesStr}")
544+
542545
out.puts(s"func (v $fullEnumNameStr) isDefined() bool {")
543546
out.inc
544-
out.puts(s"_, ok := values_$fullEnumNameStr[v]")
547+
out.puts(s"_, ok := valueNames_$fullEnumNameStr[v]")
545548
out.puts("return ok")
546549
out.dec
547550
out.puts("}")
551+
552+
out.puts(s"func (v $fullEnumNameStr) IsA$fullEnumNameStr() bool {")
553+
out.inc
554+
out.puts(s"return v.isDefined()")
555+
out.dec
556+
out.puts("}")
557+
558+
out.puts(s"func (v $fullEnumNameStr) String() string {")
559+
out.inc
560+
out.puts(s"name, ok := valueNames_$fullEnumNameStr[v]")
561+
out.puts("if ok {")
562+
out.inc
563+
out.puts("return name")
564+
out.dec
565+
out.puts("}")
566+
out.puts(s"return fmt.Sprintf(\"${fullEnumNameStr}(%d)\", v)")
567+
out.dec
568+
out.puts("}")
569+
570+
out.puts(s"func ${fullEnumNameStr}Strings() []string {")
571+
out.inc
572+
out.puts(s"strings := make([]string, 0, len(valueNames_$fullEnumNameStr))")
573+
out.puts(s"for _, name := range valueNames_$fullEnumNameStr {")
574+
out.inc
575+
out.puts("strings = append(strings, name)")
576+
out.dec
577+
out.puts("}")
578+
out.puts("return strings")
579+
out.dec
580+
out.puts("}")
581+
582+
out.puts(s"func ${fullEnumNameStr}Values() []${fullEnumNameStr} {")
583+
out.inc
584+
out.puts(s"values := make([]${fullEnumNameStr}, 0, len(valueNames_$fullEnumNameStr))")
585+
out.puts(s"for value := range valueNames_$fullEnumNameStr {")
586+
out.inc
587+
out.puts("values = append(values, value)")
588+
out.dec
589+
out.puts("}")
590+
out.puts("return values")
591+
out.dec
592+
out.puts("}")
548593
}
549594

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

0 commit comments

Comments
 (0)