diff --git a/jen/dict.go b/jen/dict.go index 260cf92..a3d98a1 100644 --- a/jen/dict.go +++ b/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 { diff --git a/jen/examples_test.go b/jen/examples_test.go index f1d825a..4990cf9 100644 --- a/jen/examples_test.go +++ b/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"), diff --git a/jen/ordered_dict.go b/jen/ordered_dict.go new file mode 100644 index 0000000..dd74322 --- /dev/null +++ b/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 +}