From fa7370477d68d08915fba8b174744c3c9ed250e1 Mon Sep 17 00:00:00 2001
From: wyc <wyc@fastmail.fm>
Date: Fri, 4 Nov 2016 18:03:45 -0400
Subject: [PATCH] Add Optional bson Support

Adds a new flag, `-bson`, which will also generate `MarshalBSON` and
`UnmarshalBSON` commands that satisfy interfaces specified in
gopkg.in/mgo.v2/bson.
---
 jsonenums.go | 15 +++++++++++++++
 template.go  | 31 +++++++++++++++++++++----------
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/jsonenums.go b/jsonenums.go
index 60346ee..7dcbfc4 100644
--- a/jsonenums.go
+++ b/jsonenums.go
@@ -82,6 +82,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() {
@@ -109,14 +112,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 7c14f0c..f741348 100644
--- a/template.go
+++ b/template.go
@@ -6,17 +6,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}}
 
@@ -42,22 +49,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]
@@ -68,5 +77,7 @@ func (r *{{$typename}}) UnmarshalJSON(data []byte) error {
     return nil
 }
 
+{{end}}
+
 {{end}}
 `))