Skip to content

Instantly share code, notes, and snippets.

@warpfork
Created May 2, 2019 21:50
Show Gist options
  • Save warpfork/1e5b72f011c14885fc358287edadb2d3 to your computer and use it in GitHub Desktop.
Save warpfork/1e5b72f011c14885fc358287edadb2d3 to your computer and use it in GitHub Desktop.
a code generator for enums for golang that i should probably push somewhere serious
package main
/*
This is all explicitly designed to be compatible with BurntSushi's sumtype checker:
https://github.com/BurntSushi/go-sumtype
Concretely, that means we 1) generate an interface and 2) use the unexported method
trick to make it 'sealed'.
All enum values get their own (empty) struct type.
This is so values cannot be coerced to them; and all values of that type, all being
zero, are all equal (and also consume no space, if you really go deep into Go;
the compiler automagically singletonizes them).
Philosophically, sum types and enums are quite close; enums simply have no further
attached information. The BurntSushi checker for exhaustiveness of sum type switches
actually supports much *more* power than we need for enums; our approach to making
an empty struct type for each enum value simply makes it trivially compatible.
Your user code should refer to the interface whenever refering to the sum type...
except in serializable/marshallable message struct types, in which case an additional
wrapper struct is provided, which will check incoming data matches known ranges.
When switching on enum values, prefer to use type switches. These are faster,
and also are checkable by the BurntSushi utility.
(Value switches will also work, but we don't have a tool for static checking them.)
The {Enum}Values list and {Enum}Lookup map are both mutatable, simply because in Go
there's no way to be immutable. You should treat them as immutable nonetheless.
They are there for convenient ranging operations. Mutating them will make your
program unpredictable and Bad.
Hypothetically, one could also build enums which simply *are* typedefs of strings;
getting the value of statically checkable enums would then also require checking
that the entire codebase never does any *casts* to those values, because the 'sealed'
trick doesn't work for something you can cast into. It also wouldn't be compatible
with the BurntSushi checker.
Some other resources on the topic:
- https://medium.com/where-do-we-go-now/stricter-go-enums-with-go-genums-2d7affc8eb0e
- https://www.reddit.com/r/golang/comments/60dik1/a_utility_for_running_exhaustiveness_checks_on/
- https://golang.org/src/go/ast/walk.go?s=1311:1342#L41
This is not the first enum generator written for go. Here are some others:
- https://github.com/abice/go-enum
- This is iota-based and so subject to casting and cannot be "closed", and
thus cannot be used with the sumType exhaustiveness checker;
- I'll leave the reader to form their own opinions about that "_ColorName" in
the examples...
- https://github.com/gdm85/go-genums
- This somehow comes with a constraint of "Do not use the enum types without
calling their New() method", which seems like a bit of a footgun.
To consider: we could make all of the types have names suffixed with
"T" (c-style), and export consts^W vars in order to save our users the agony
of trailing "{}" at use sites. However, this would be A) another "var"
situation where you can assign six to seven, so to speak; and B) cause one
to use "TheEnum_ValueT" in type switches.
*/
import (
"bytes"
"fmt"
"strings"
"text/template"
)
// We operate in one of two modes:
// Either an a json file with our full configuration, or,
// a series of flags for each value.
//
// go-genenums pkgname:MyEnum -- alpha beta gamma CodeName:delta Overriden:omega
//
// Colons can be used in this way because ... no they can't actually, unless they're first and then codename is first
//
// you know what, screw it, you can edit the main method.
//
func main() {
// !!!! LISTEN UP !!!!
// YOU CONFIGURE THIS THANG RIGHT HERE
// BY EDITING THIS COMMENT BLOCK, YEAH
// BECAUSE I WAS TOO LAZY TO STRAP A REAL CLI OR ANYTHING ELSE ON THIS THING
// emit(enum{
// PackageName: "pkg",
// TypeName: "myEnum",
// VarsPrefix: "myEnum_",
// Values: []enumValue{
// {Value: "alpha"}, {Value: "beta"}, {Value: "gamma"}, {Value: "delta"}, {Value: "omega"},
// },
// }.defaults())
}
type enum struct {
PackageName string // mostly not used except for file header itself.
TypeName string // type name. Capital to be exported, as usual.
VarsPrefix string // prefix for variables. by default, "{TypeName}_".
Values []enumValue // values. they're always strings.
}
type enumValue struct {
Value string // how both `Stringer` and `JSONMarshaller` will string this, as well as the {TypeName}Lookup map key.
CodeName string // the value name in code. default is first-letter capitalized of `values`.
}
func (en enum) defaults() enum {
evs2 := make([]enumValue, len(en.Values))
copy(evs2, en.Values)
for i := range evs2 {
if evs2[i].CodeName == "" {
evs2[i].CodeName = strings.Title(evs2[i].Value)
}
}
en.Values = evs2
return en
}
const outputTemplate = `{{$enum := .}}
package {{.PackageName}}
// *** generated with go-genenums ***
//go:generate go-genenums
import (
"encoding/json"
"fmt"
)
//go-sumtype:decl {{.TypeName}}
// {{.TypeName}} is an enum interface.
//
// The following values are the valid members of the enumeration:
//
{{range $value := .Values -}}
// {{$enum.PackageName}}.{{$enum.VarsPrefix}}{{$value.CodeName}}
{{end -}}
//
type {{.TypeName}} interface {
String() string
_sealed_{{.TypeName}}()
}
var {{.TypeName}}Values = [{{len .Values}}]{{.TypeName}}{
{{- range $value := .Values -}}
{{"\n\t"}}{{$enum.VarsPrefix}}{{$value.CodeName}}{},
{{- end}}
}
var {{.TypeName}}Lookup = map[string]{{.TypeName}}{
{{- range $value := .Values -}}
{{- /* Note that this will change on gofmt due to alignment. */ -}}
{{"\n\t"}}"{{$value.Value}}": {{$enum.VarsPrefix}}{{$value.CodeName}}{},
{{- end}}
}
func Reify{{.TypeName|toTitle}}(s string) ({{.TypeName}}, error) {
if val, ok := {{.TypeName}}Lookup[s]; ok {
return val, nil
} else {
return nil, fmt.Errorf("{{.PackageName}}.{{.TypeName}}: %q is not a known enum value", s)
}
}
{{range $value := .Values -}}
type {{$enum.VarsPrefix}}{{$value.CodeName}} struct{}
func (x {{$enum.VarsPrefix}}{{$value.CodeName}}) _sealed_{{$enum.TypeName}}() {}
func (x {{$enum.VarsPrefix}}{{$value.CodeName}}) String() string {
return "{{$value.Value}}"
}
func (x {{$enum.VarsPrefix}}{{$value.CodeName}}) GoString() string {
return "{{$enum.VarsPrefix}}{{$value.CodeName}}"
}
{{end -}}
// {{.TypeName}}Marshaller is a helper for use in structs which will be marshalled
// or unmarshalled with stdlib json marshalling.
//
// (We wouldn't need this if it was possible to attach marshaller functions to
// interfaces, or attach useful metadata to types via tags, but alas, Golang.)
type {{.TypeName}}Marshaller struct {
Value {{.TypeName}}
}
func (m {{.TypeName}}Marshaller) MarshalJSON() ([]byte, error) {
return []byte(fmt.Sprintf("%q", m.Value.String())), nil
}
func (m *{{.TypeName}}Marshaller) UnmarshalJSON(data []byte) error {
var name string
if err := json.Unmarshal(data, &name); err != nil {
return fmt.Errorf("invalid data for unmarshalling {{.PackageName}}.{{.TypeName}}, expected string: %s", err)
}
if val, err := Reify{{.TypeName|toTitle}}(name); err != nil {
return fmt.Errorf("invalid data for unmarshalling %s", err)
} else {
m.Value = val
}
return nil
}
`
func emit(en enum) {
fmt.Println(tmpl(outputTemplate, en))
}
func tmpl(tmpl string, obj interface{}) string {
t := template.Must(template.New("").Funcs(template.FuncMap{
"toTitle": strings.Title,
}).Parse(tmpl))
var buf bytes.Buffer
if err := t.Execute(&buf, obj); err != nil {
panic(err)
}
return buf.String()
}
@ghostsquad
Copy link

So when are you going to make this a real CLI?

@ghostsquad
Copy link

More work is needed, but here's how I did it locally:

func main() {
	var (
		out = flag.String("out", "", "file to save output")
		pkgName = flag.String("pkg", "", "package name for generated files")
		typeName = flag.String("type", "", "type name")
		values = flag.String("values", "", "comma separated list of values")
	)

	flag.Parse()

	fmt.Println("output file:", *out)

	var enumValues []enumValue

	for _, v := range strings.Split(*values, ",") {
		enumValues = append(enumValues, enumValue{ Value: v })
	}

	lf := &LazyFile{FileName: *out}
	defer lf.Close()

	emit(enum{
		PackageName: *pkgName,
		TypeName: *typeName,
		VarsPrefix: "",
		Values: enumValues,
	}.defaults(), lf)
}

also add these:

// LazyFile is an io.WriteCloser which defers creation of the file it is supposed to write in
// till the first call to its write function in order to prevent creation of file, if no write
// is supposed to happen.
type LazyFile struct {
	// FileName is path to the file to which genny will write.
	FileName string
	file     *os.File
}

// Close closes the file if it is created. Returns nil if no file is created.
func (lw *LazyFile) Close() error {
	if lw.file != nil {
		return lw.file.Close()
	}
	return nil
}

// Write writes to the specified file and creates the file first time it is called.
func (lw *LazyFile) Write(p []byte) (int, error) {
	if lw.file == nil {
		err := os.MkdirAll(path.Dir(lw.FileName), 0755)
		if err != nil {
			return 0, err
		}
		lw.file, err = os.Create(lw.FileName)
		if err != nil {
			return 0, err
		}
	}
	return lw.file.Write(p)
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment