sortable_yaml.go

package yaml

import (
	"sort"

	"gopkg.in/yaml.v2"
	sigsyaml "sigs.k8s.io/yaml"
)

type orderFunc func(a, b string) bool

type customString struct {
	s []string
	f orderFunc
}

func (a customString) Len() int           { return len(a.s) }
func (a customString) Less(i, j int) bool { return a.f(a.s[i], a.s[j]) }
func (a customString) Swap(i, j int)      { a.s[i], a.s[j] = a.s[j], a.s[i] }

func OrderedMarshal(o interface{}, f orderFunc) ([]byte, error) {
	rawYaml, err := sigsyaml.Marshal(o)
	if err != nil {
		return nil, err
	}
	yamlMap := map[string]interface{}{}
	err = sigsyaml.Unmarshal(rawYaml, &yamlMap)

	keys := []string{}
	for k := range yamlMap {
		keys = append(keys, k)
	}
	sort.Sort(customString{s: keys, f: f})

	orderedData := yaml.MapSlice{}
	for _, ok := range keys {
		orderedData = append(orderedData, yaml.MapItem{Key: ok, Value: yamlMap[ok]})
	}
	return yaml.Marshal(orderedData)
}

sortable_yaml_test.go

package yaml

import (
	"testing"
)

type testStruct struct {
	Name   string `json:"name,omitempty"`
	Colour int    `json:"colour,omitempty"`
	Height int    `json:"height,omitempty"`
}

func TestOrderedYAML(t *testing.T) {
	b, err := OrderedMarshal(&testStruct{Name: "testing", Colour: 20, Height: 2}, func(a, b string) bool {
		if a == "name" {
			return true
		}
		if b == "name" {
			return false
		}
		return a < b
	})
	if err != nil {
		t.Fatal(err)
	}

	want := "name: testing\ncolour: 20\nheight: 2\n"
	if want != string(b) {
		t.Fatalf("\ngot:\n%s\nwant:\n%s", string(b), want)
	}
}

func TestOrderedYAMLWithCustomOrdering(t *testing.T) {
	desirableFunc := func(d []string) orderFunc {
		return func(a, b string) bool {
			ai := stringIndex(a, d)
			bi := stringIndex(b, d)
			if ai > -1 && bi > -1 {
				return ai < bi
			}
			if ai > -1 {
				return true
			}
			return false

		}
	}([]string{"name", "height"})

	b, err := OrderedMarshal(&testStruct{Name: "testing", Colour: 20, Height: 4}, desirableFunc)
	if err != nil {
		t.Fatal(err)
	}

	want := "name: testing\nheight: 4\ncolour: 20\n"
	if want != string(b) {
		t.Fatalf("\ngot:\n%s\nwant:\n%s", string(b), want)
	}
}

func stringIndex(s string, ss []string) int {
	for i, v := range ss {
		if s == v {
			return i
		}
	}
	return -1
}