paralus/cmd/generate-enum/main.go

package main

import (
	"fmt"
	"os"
	"strings"
	"text/template"
)

const tmpl = `
// Code generated by go generate; DO NOT EDIT.
package {{ .PackageName }}

import (
	driver "database/sql/driver"
	bytes "bytes"
	"fmt"
)

// Scan converts database string to {{ .EnumName }}
func (e *{{ .EnumName }}) Scan(value interface{}) error {
	s := value.([]byte)
	*e = {{ .EnumName }}({{ .EnumName }}_value[string(s)])
	return nil
}

// Value converts {{ .EnumName }} into database string
func (e {{ .EnumName }}) Value() (driver.Value, error) {
	return {{ .EnumName }}_name[int32(e)], nil
}

// MarshalJSON converts {{ .EnumName }} to JSON
func (e {{ .EnumName }}) MarshalJSON() ([]byte, error) {
	buffer := bytes.NewBufferString("\"")
	buffer.WriteString(e.String())
	buffer.WriteString("\"")
	return buffer.Bytes(), nil
}

// UnmarshalJSON converts {{ .EnumName }} from JSON
func (e *{{ .EnumName }}) UnmarshalJSON(b []byte) error {
	if b != nil {
		*e = {{ .EnumName }}({{ .EnumName }}_value[string(b[1:len(b)-1])])
	}
	return nil
}

// MarshalYAML implements the yaml.Marshaler interface 
func (e {{ .EnumName }}) MarshalYAML() (interface{}, error) {
	return {{ .EnumName }}_name[int32(e)], nil
}

// UnmarshalYAML implements the yaml.Unmarshaler interface
func (e *{{ .EnumName }}) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var name string
	if err := unmarshal(&name); err != nil {
		return err
	}

	value, ok := {{ .EnumName }}_value[name]
	if !ok {
		return fmt.Errorf("invalid {{ .EnumName }}: %s", name)
	}

	*e = {{ .EnumName }}(value)
	return nil
}

// implement proto enum interface
func (e {{ .EnumName }}) IsEnum() {
}

`

type TemplateData struct {
	PackageName string
	EnumName    string
}

func main() {
	if len(os.Args) != 3 {
		fmt.Println("expected generate-enum <enum_type> <gen_file_dir>")
		fmt.Println("got", os.Args)
		os.Exit(1)
	}
	s := strings.Split(os.Args[2], "/")
	packageName := s[len(s)-1]

	fName := fmt.Sprintf("%s/%s.enum.go", os.Args[2], strings.ToLower(os.Args[1]))

	f, err := os.Create(fName)
	if err != nil {
		fmt.Println("unable to create file ", err)
		os.Exit(1)
	}
	defer f.Close()
	pkgTemplate := template.Must(template.New("").Parse(tmpl))
	err = pkgTemplate.Execute(f, TemplateData{
		PackageName: packageName,
		EnumName:    os.Args[1],
	})
	if err != nil {
		fmt.Println("render template ", err)
		os.Exit(1)
	}
}