From ba6c96412b841f652a2a7ff28ca55b8b30141d6e Mon Sep 17 00:00:00 2001 From: Alexey Khit Date: Tue, 25 Jul 2023 18:05:50 +0300 Subject: [PATCH] Add YAML pkg with Patch function --- pkg/yaml/yaml.go | 185 ++++++++++++++++++++++++++++++++++++++++++ pkg/yaml/yaml_test.go | 85 +++++++++++++++++++ 2 files changed, 270 insertions(+) create mode 100644 pkg/yaml/yaml.go create mode 100644 pkg/yaml/yaml_test.go diff --git a/pkg/yaml/yaml.go b/pkg/yaml/yaml.go new file mode 100644 index 00000000..7d9d6345 --- /dev/null +++ b/pkg/yaml/yaml.go @@ -0,0 +1,185 @@ +package yaml + +import ( + "bytes" + "errors" + + "gopkg.in/yaml.v3" +) + +func Encode(v any, indent int) ([]byte, error) { + b := bytes.NewBuffer(nil) + e := yaml.NewEncoder(b) + e.SetIndent(indent) + + if err := e.Encode(v); err != nil { + return nil, err + } + + return b.Bytes(), nil +} + +// Patch - change key/value pair in YAML file without break formatting +func Patch(src []byte, key string, value any, path ...string) ([]byte, error) { + nodeParent, err := FindParent(src, path...) + if err != nil { + return nil, err + } + + var dst []byte + + if nodeParent != nil { + dst, err = AddOrReplace(src, key, value, nodeParent) + } else { + dst, err = AddToEnd(src, key, value, path...) + } + + if err = yaml.Unmarshal(dst, map[string]any{}); err != nil { + return nil, err + } + + return dst, nil +} + +// FindParent - return YAML Node from path of keys (tree) +func FindParent(src []byte, path ...string) (*yaml.Node, error) { + if len(src) == 0 { + return nil, nil + } + + var root yaml.Node + if err := yaml.Unmarshal(src, &root); err != nil { + return nil, err + } + + if root.Content == nil { + return nil, nil + } + + parent := root.Content[0] // yaml.DocumentNode + for _, name := range path { + if parent == nil { + break + } + _, parent = FindChild(parent, name) + } + return parent, nil +} + +// FindChild - search and return YAML key/value pair for current Node +func FindChild(node *yaml.Node, name string) (key, value *yaml.Node) { + for i, child := range node.Content { + if child.Value != name { + continue + } + return child, node.Content[i+1] + } + + return nil, nil +} + +func FirstChild(node *yaml.Node) *yaml.Node { + if node.Content == nil { + return node + } + return node.Content[0] +} + +func LastChild(node *yaml.Node) *yaml.Node { + if node.Content == nil { + return node + } + return node.Content[len(node.Content)-1] +} + +func AddOrReplace(src []byte, key string, value any, nodeParent *yaml.Node) ([]byte, error) { + v := map[string]any{key: value} + put, err := Encode(v, 2) + if err != nil { + return nil, err + } + + if nodeKey, nodeValue := FindChild(nodeParent, key); nodeKey != nil { + put = AddIndent(put, nodeKey.Column-1) + + i0 := LineOffset(src, nodeKey.Line) + i1 := LineOffset(src, LastChild(nodeValue).Line+1) + + dst := make([]byte, 0, len(src)+len(put)) + dst = append(dst, src[:i0]...) + if value != nil { + dst = append(dst, put...) + } + return append(dst, src[i1:]...), nil + } + + put = AddIndent(put, FirstChild(nodeParent).Column-1) + + i := LineOffset(src, LastChild(nodeParent).Line+1) + + dst := make([]byte, 0, len(src)+len(put)) + dst = append(dst, src[:i]...) + if value != nil { + dst = append(dst, put...) + } + return append(dst, src[i:]...), nil +} + +func AddToEnd(src []byte, key string, value any, path ...string) ([]byte, error) { + if len(path) > 1 || value == nil { + return nil, errors.New("config: path not exist") + } + + v := map[string]map[string]any{ + path[0]: {key: value}, + } + put, err := Encode(v, 2) + if err != nil { + return nil, err + } + + dst := make([]byte, 0, len(src)+len(put)+10) + dst = append(dst, src...) + if l := len(src); l > 0 && src[l-1] != '\n' { + dst = append(dst, '\n') + } + return append(dst, put...), nil +} + +func AddPrefix(src, pre []byte) (dst []byte) { + for len(src) > 0 { + dst = append(dst, pre...) + i := bytes.IndexByte(src, '\n') + 1 + if i == 0 { + dst = append(dst, src...) + break + } + dst = append(dst, src[:i]...) + src = src[i:] + } + + return +} + +func AddIndent(src []byte, indent int) (dst []byte) { + pre := make([]byte, indent) + for i := 0; i < indent; i++ { + pre[i] = ' ' + } + return AddPrefix(src, pre) +} + +func LineOffset(b []byte, line int) (offset int) { + for l := 1; ; l++ { + if l == line { + return offset + } + + i := bytes.IndexByte(b[offset:], '\n') + 1 + if i == 0 { + break + } + offset += i + } + return -1 +} diff --git a/pkg/yaml/yaml_test.go b/pkg/yaml/yaml_test.go new file mode 100644 index 00000000..7547e19b --- /dev/null +++ b/pkg/yaml/yaml_test.go @@ -0,0 +1,85 @@ +package yaml + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestPatch(t *testing.T) { + b := []byte(`# prefix`) + + b, err := Patch(b, "camera1", "url1", "streams") + require.Nil(t, err) + + require.Equal(t, `# prefix +streams: + camera1: url1 +`, string(b)) + + b, err = Patch(b, "camera2", []string{"url2", "url3"}, "streams") + require.Nil(t, err) + + require.Equal(t, `# prefix +streams: + camera1: url1 + camera2: + - url2 + - url3 +`, string(b)) + + b, err = Patch(b, "camera1", "url4", "streams") + require.Nil(t, err) + + require.Equal(t, `# prefix +streams: + camera1: url4 + camera2: + - url2 + - url3 +`, string(b)) + + b, err = Patch(b, "camera2", "url5", "streams") + require.Nil(t, err) + + require.Equal(t, `# prefix +streams: + camera1: url4 + camera2: url5 +`, string(b)) + + b, err = Patch(b, "camera1", nil, "streams") + require.Nil(t, err) + + require.Equal(t, `# prefix +streams: + camera2: url5 +`, string(b)) +} + +func TestPatchParings(t *testing.T) { + b := []byte(`homekit: + camera1: + pin: 123-45-678 +streams: + camera1: url1 +`) + + pairings := map[string]string{ + "client1": "public1", + "client2": "public2", + } + + b, err := Patch(b, "pairings", pairings, "homekit", "camera1") + require.Nil(t, err) + + require.Equal(t, `homekit: + camera1: + pin: 123-45-678 + pairings: + client1: public1 + client2: public2 +streams: + camera1: url1 +`, string(b)) +}