Add OrderedDict. Refactor Dict.render to utilize OrderedDict

jen/dict.go

@@ -19,51 +19,32 @@ func DictFunc(f func(Dict)) Dict {
 }
 
 func (d Dict) render(f *File, w io.Writer, s *Statement) error {
-	first := true
-	// must order keys to ensure repeatable source
-	type kv struct {
-		k Code
-		v Code
-	}
-	lookup := map[string]kv{}
-	keys := []string{}
+	lookup := make(map[string]kv, len(d))
+	keys := make([]string, 0, len(d))
+
+	buf := &bytes.Buffer{}
 	for k, v := range d {
 		if k.isNull(f) || v.isNull(f) {
 			continue
 		}
-		buf := &bytes.Buffer{}
 		if err := k.render(f, buf, nil); err != nil {
 			return err
 		}
 		keys = append(keys, buf.String())
 		lookup[buf.String()] = kv{k: k, v: v}
+		buf.Reset()
 	}
+
+	// must order keys to ensure repeatable source
 	sort.Strings(keys)
+
+	ordered := make([]kv, 0, len(keys))
 	for _, key := range keys {
-		k := lookup[key].k
-		v := lookup[key].v
-		if first && len(keys) > 1 {
-			if _, err := w.Write([]byte("\n")); err != nil {
-				return err
-			}
-			first = false
-		}
-		if err := k.render(f, w, nil); err != nil {
-			return err
-		}
-		if _, err := w.Write([]byte(":")); err != nil {
-			return err
-		}
-		if err := v.render(f, w, nil); err != nil {
-			return err
-		}
-		if len(keys) > 1 {
-			if _, err := w.Write([]byte(",\n")); err != nil {
-				return err
-			}
-		}
+		ordered = append(ordered, kv{k: lookup[key].k, v: lookup[key].v})
 	}
-	return nil
+
+	dict := &OrderedDict{items: ordered}
+	return dict.render(f, w, s)
 }
 
 func (d Dict) isNull(f *File) bool {

jen/examples_test.go

@@ -1577,6 +1577,26 @@ func ExampleDictFunc() {
 	// }
 }
 
+func ExampleOrderedDictFunc() {
+	c := Id("a").Op(":=").Map(String()).String().Values(OrderedDictFunc(func(o *OrderedDict) {
+		o.Add(Lit("c"), Lit("d"))
+		o.Add(Lit("a"), Lit("b"))
+	}))
+	fmt.Printf("%#v", c)
+	// Output:
+	// a := map[string]string{
+	// 	"c": "d",
+	// 	"a": "b",
+	// }
+}
+
+func ExampleOrderedDictFunc_nil() {
+	c := Id("a").Op(":=").Map(String()).String().Values(OrderedDictFunc(func(o *OrderedDict) {}))
+	fmt.Printf("%#v", c)
+	// Output:
+	// a := map[string]string{}
+}
+
 func ExampleDefs() {
 	c := Const().Defs(
 		Id("a").Op("=").Lit("a"),

jen/ordered_dict.go

@@ -0,0 +1,69 @@
+package jen
+
+import "io"
+
+// OrderedDictFunc executes a func(*OrderedDict) to generate the value. Use OrderedDict.Add to append key/value pairs.
+// Use with Values for map or composite literals.
+func OrderedDictFunc(f func(*OrderedDict)) *OrderedDict {
+	o := &OrderedDict{}
+	f(o)
+	return o
+}
+
+// OrderedDict renders as key/value pairs. Use with Values for map or composite
+// literals.
+type OrderedDict struct {
+	items []kv
+}
+
+type kv struct {
+	k Code
+	v Code
+}
+
+// Add appends a key/value pairs.
+func (o *OrderedDict) Add(key Code, value Code) *OrderedDict {
+	o.items = append(o.items, kv{k: key, v: value})
+	return o
+}
+
+func (o *OrderedDict) render(f *File, w io.Writer, _ *Statement) error {
+	first := true
+	for _, item := range o.items {
+		if first && len(o.items) > 1 {
+			if _, err := w.Write([]byte("\n")); err != nil {
+				return err
+			}
+			first = false
+		}
+		if err := item.k.render(f, w, nil); err != nil {
+			return err
+		}
+		if _, err := w.Write([]byte(":")); err != nil {
+			return err
+		}
+		if err := item.v.render(f, w, nil); err != nil {
+			return err
+		}
+		if len(o.items) > 1 {
+			if _, err := w.Write([]byte(",\n")); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (o *OrderedDict) isNull(f *File) bool {
+	if len(o.items) == 0 {
+		return true
+	}
+	for _, item := range o.items {
+		if !item.k.isNull(f) && !item.v.isNull(f) {
+			// if any of the key/value pairs are both not null, the OrderedDict is not
+			// null
+			return false
+		}
+	}
+	return true
+}