123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204 |
- package yaml
- import (
- "bytes"
- "errors"
- "gopkg.in/yaml.v3"
- )
- func Unmarshal(in []byte, out interface{}) (err error) {
- return yaml.Unmarshal(in, out)
- }
- 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 LastChild(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)
- if i1 < 0 { // no new line on the end of file
- if value != nil {
- return append(src[:i0], put...), nil
- }
- return src[:i0], nil
- }
- 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)
- if i < 0 { // no new line on the end of file
- src = append(src, '\n')
- if value != nil {
- src = append(src, put...)
- }
- return src, nil
- }
- 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
- }
|