-
Notifications
You must be signed in to change notification settings - Fork 46
/
yaml_map.go
204 lines (185 loc) · 4.6 KB
/
yaml_map.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
// Package yamlmap is a wrapper of gopkg.in/yaml.v3 for interacting
// with yaml data as if it were a map.
package yamlmap
import (
"errors"
"gopkg.in/yaml.v3"
)
const (
modified = "modifed"
)
type Map struct {
*yaml.Node
}
var ErrNotFound = errors.New("not found")
var ErrInvalidYaml = errors.New("invalid yaml")
var ErrInvalidFormat = errors.New("invalid format")
func StringValue(value string) *Map {
return &Map{&yaml.Node{
Kind: yaml.ScalarNode,
Tag: "!!str",
Value: value,
}}
}
func MapValue() *Map {
return &Map{&yaml.Node{
Kind: yaml.MappingNode,
Tag: "!!map",
}}
}
func Unmarshal(data []byte) (*Map, error) {
var root yaml.Node
err := yaml.Unmarshal(data, &root)
if err != nil {
return nil, ErrInvalidYaml
}
if len(root.Content) == 0 || root.Content[0].Kind != yaml.MappingNode {
return nil, ErrInvalidFormat
}
return &Map{root.Content[0]}, nil
}
func Marshal(m *Map) ([]byte, error) {
return yaml.Marshal(m.Node)
}
func (m *Map) AddEntry(key string, value *Map) {
keyNode := &yaml.Node{
Kind: yaml.ScalarNode,
Tag: "!!str",
Value: key,
}
m.Content = append(m.Content, keyNode, value.Node)
m.SetModified()
}
func (m *Map) Empty() bool {
return m.Content == nil || len(m.Content) == 0
}
func (m *Map) FindEntry(key string) (*Map, error) {
// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
// When iterating over the content slice we only want to compare the keys of the yamlMap.
for i, v := range m.Content {
if i%2 != 0 {
continue
}
if v.Value == key {
if i+1 < len(m.Content) {
return &Map{m.Content[i+1]}, nil
}
}
}
return nil, ErrNotFound
}
func (m *Map) Keys() []string {
// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
// When iterating over the content slice we only want to select the keys of the yamlMap.
keys := []string{}
for i, v := range m.Content {
if i%2 != 0 {
continue
}
keys = append(keys, v.Value)
}
return keys
}
func (m *Map) RemoveEntry(key string) error {
// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
// When iterating over the content slice we only want to compare the keys of the yamlMap.
// If we find they key to remove, remove the key and its value from the content slice.
found, skipNext := false, false
newContent := []*yaml.Node{}
for i, v := range m.Content {
if skipNext {
skipNext = false
continue
}
if i%2 != 0 || v.Value != key {
newContent = append(newContent, v)
} else {
found = true
skipNext = true
m.SetModified()
}
}
if !found {
return ErrNotFound
}
m.Content = newContent
return nil
}
func (m *Map) SetEntry(key string, value *Map) {
// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
// When iterating over the content slice we only want to compare the keys of the yamlMap.
// If we find they key to set, set the next item in the content slice to the new value.
m.SetModified()
for i, v := range m.Content {
if i%2 != 0 || v.Value != key {
continue
}
if v.Value == key {
if i+1 < len(m.Content) {
m.Content[i+1] = value.Node
return
}
}
}
m.AddEntry(key, value)
}
// Note: This is a hack to introduce the concept of modified/unmodified
// on top of gopkg.in/yaml.v3. This works by setting the Value property
// of a MappingNode to a specific value and then later checking if the
// node's Value property is that specific value. When a MappingNode gets
// output as a string the Value property is not used, thus changing it
// has no impact for our purposes.
func (m *Map) SetModified() {
// Can not mark a non-mapping node as modified
if m.Node.Kind != yaml.MappingNode && m.Node.Tag == "!!null" {
m.Node.Kind = yaml.MappingNode
m.Node.Tag = "!!map"
}
if m.Node.Kind == yaml.MappingNode {
m.Node.Value = modified
}
}
// Traverse map using BFS to set all nodes as unmodified.
func (m *Map) SetUnmodified() {
i := 0
queue := []*yaml.Node{m.Node}
for {
if i > (len(queue) - 1) {
break
}
q := queue[i]
i = i + 1
if q.Kind != yaml.MappingNode {
continue
}
q.Value = ""
queue = append(queue, q.Content...)
}
}
// Traverse map using BFS to searach for any nodes that have been modified.
func (m *Map) IsModified() bool {
i := 0
queue := []*yaml.Node{m.Node}
for {
if i > (len(queue) - 1) {
break
}
q := queue[i]
i = i + 1
if q.Kind != yaml.MappingNode {
continue
}
if q.Value == modified {
return true
}
queue = append(queue, q.Content...)
}
return false
}
func (m *Map) String() string {
data, err := Marshal(m)
if err != nil {
return ""
}
return string(data)
}