diff --git a/jsonenums.go b/jsonenums.go index 7d26291..94c50ad 100644 --- a/jsonenums.go +++ b/jsonenums.go @@ -91,6 +91,9 @@ var ( typeNames = flag.String("type", "", "comma-separated list of type names; must be set") outputPrefix = flag.String("prefix", "", "prefix to be added to the output file") outputSuffix = flag.String("suffix", "_jsonenums", "suffix to be added to the output file") + + // BSON Support + bsonMode = flag.Bool("bson", false, "enable BSON-mode and generate BSON output in addition to JSON") ) func main() { @@ -118,14 +121,26 @@ func main() { log.Fatalf("parsing package: %v", err) } + funcPrefixes := []string{"JSON"} + imports := []string{"encoding/json"} + if *bsonMode { + funcPrefixes = append(funcPrefixes, "BSON") + imports = append(imports, "gopkg.in/mgo.v2/bson") + } var analysis = struct { Command string PackageName string TypesAndValues map[string][]string + + // ["JSON", "BSON"] + FuncPrefixes []string + Imports []string }{ Command: strings.Join(os.Args[1:], " "), PackageName: pkg.Name, TypesAndValues: make(map[string][]string), + FuncPrefixes: funcPrefixes, + Imports: imports, } // Run generate for each type. diff --git a/template.go b/template.go index 332c2a3..76aa517 100644 --- a/template.go +++ b/template.go @@ -15,17 +15,24 @@ package main -import "text/template" +import ( + "strings" + "text/template" +) -var generatedTmpl = template.Must(template.New("generated").Parse(` +var generatedTmpl = template.Must(template.New("generated"). + Funcs(template.FuncMap{"toLower": strings.ToLower}).Parse(` // generated by jsonenums {{.Command}}; DO NOT EDIT package {{.PackageName}} import ( - "encoding/json" +{{range .Imports}} + "{{.}}" +{{end}} "fmt" ) +{{$funcPrefixes := .FuncPrefixes}} {{range $typename, $values := .TypesAndValues}} @@ -51,22 +58,24 @@ func init() { } } -// MarshalJSON is generated so {{$typename}} satisfies json.Marshaler. -func (r {{$typename}}) MarshalJSON() ([]byte, error) { +{{ range $_, $funcPrefix := $funcPrefixes}} + +// Marshal{{$funcPrefix}} is generated so {{$typename}} satisfies {{$funcPrefix | toLower}}.Marshaler. +func (r {{$typename}}) Marshal{{$funcPrefix}}() ([]byte, error) { if s, ok := interface{}(r).(fmt.Stringer); ok { - return json.Marshal(s.String()) + return {{$funcPrefix | toLower}}.Marshal(s.String()) } s, ok := _{{$typename}}ValueToName[r] if !ok { return nil, fmt.Errorf("invalid {{$typename}}: %d", r) } - return json.Marshal(s) + return {{$funcPrefix | toLower}}.Marshal(s) } -// UnmarshalJSON is generated so {{$typename}} satisfies json.Unmarshaler. -func (r *{{$typename}}) UnmarshalJSON(data []byte) error { +// Unmarshal{{$funcPrefix}} is generated so {{$typename}} satisfies {{$funcPrefix | toLower}}.Unmarshaler. +func (r *{{$typename}}) Unmarshal{{$funcPrefix}}(data []byte) error { var s string - if err := json.Unmarshal(data, &s); err != nil { + if err := {{$funcPrefix | toLower}}.Unmarshal(data, &s); err != nil { return fmt.Errorf("{{$typename}} should be a string, got %s", data) } v, ok := _{{$typename}}NameToValue[s] @@ -77,5 +86,7 @@ func (r *{{$typename}}) UnmarshalJSON(data []byte) error { return nil } +{{end}} + {{end}} `))