forked from mikefarah/yq
-
Notifications
You must be signed in to change notification settings - Fork 0
/
encoder.go
173 lines (161 loc) · 3.77 KB
/
encoder.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
package yqlib
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
yaml "gopkg.in/yaml.v3"
)
type Encoder interface {
Encode(writer io.Writer, node *yaml.Node) error
PrintDocumentSeparator(writer io.Writer) error
PrintLeadingContent(writer io.Writer, content string) error
CanHandleAliases() bool
}
// orderedMap allows to marshal and unmarshal JSON and YAML values keeping the
// order of keys and values in a map or an object.
type orderedMap struct {
// if this is an object, kv != nil. If this is not an object, kv == nil.
kv []orderedMapKV
altVal interface{}
}
type orderedMapKV struct {
K string
V orderedMap
}
func (o *orderedMap) UnmarshalJSON(data []byte) error {
switch data[0] {
case '{':
// initialise so that even if the object is empty it is not nil
o.kv = []orderedMapKV{}
// create decoder
dec := json.NewDecoder(bytes.NewReader(data))
_, err := dec.Token() // open object
if err != nil {
return err
}
// cycle through k/v
var tok json.Token
for tok, err = dec.Token(); !errors.Is(err, io.EOF); tok, err = dec.Token() {
// we can expect two types: string or Delim. Delim automatically means
// that it is the closing bracket of the object, whereas string means
// that there is another key.
if _, ok := tok.(json.Delim); ok {
break
}
kv := orderedMapKV{
K: tok.(string),
}
if err := dec.Decode(&kv.V); err != nil {
return err
}
o.kv = append(o.kv, kv)
}
// unexpected error
if err != nil && !errors.Is(err, io.EOF) {
return err
}
return nil
case '[':
var arr []orderedMap
return json.Unmarshal(data, &arr)
}
return json.Unmarshal(data, &o.altVal)
}
func (o orderedMap) MarshalJSON() ([]byte, error) {
buf := new(bytes.Buffer)
enc := json.NewEncoder(buf)
enc.SetEscapeHTML(false) // do not escape html chars e.g. &, <, >
if o.kv == nil {
if err := enc.Encode(o.altVal); err != nil {
return nil, err
}
return buf.Bytes(), nil
}
buf.WriteByte('{')
for idx, el := range o.kv {
if err := enc.Encode(el.K); err != nil {
return nil, err
}
buf.WriteByte(':')
if err := enc.Encode(el.V); err != nil {
return nil, err
}
if idx != len(o.kv)-1 {
buf.WriteByte(',')
}
}
buf.WriteByte('}')
return buf.Bytes(), nil
}
func (o *orderedMap) UnmarshalYAML(node *yaml.Node) error {
switch node.Kind {
case yaml.DocumentNode:
if len(node.Content) == 0 {
return nil
}
return o.UnmarshalYAML(node.Content[0])
case yaml.AliasNode:
return o.UnmarshalYAML(node.Alias)
case yaml.ScalarNode:
return node.Decode(&o.altVal)
case yaml.MappingNode:
// set kv to non-nil
o.kv = []orderedMapKV{}
for i := 0; i < len(node.Content); i += 2 {
var key string
var val orderedMap
if err := node.Content[i].Decode(&key); err != nil {
return err
}
if err := node.Content[i+1].Decode(&val); err != nil {
return err
}
o.kv = append(o.kv, orderedMapKV{
K: key,
V: val,
})
}
return nil
case yaml.SequenceNode:
// note that this has to be a pointer, so that nulls can be represented.
var res []*orderedMap
if err := node.Decode(&res); err != nil {
return err
}
o.altVal = res
o.kv = nil
return nil
case 0:
// null
o.kv = nil
o.altVal = nil
return nil
default:
return fmt.Errorf("orderedMap: invalid yaml node")
}
}
func (o *orderedMap) MarshalYAML() (interface{}, error) {
// fast path: kv is nil, use altVal
if o.kv == nil {
return o.altVal, nil
}
content := make([]*yaml.Node, 0, len(o.kv)*2)
for _, val := range o.kv {
n := new(yaml.Node)
if err := n.Encode(val.V); err != nil {
return nil, err
}
content = append(content, &yaml.Node{
Kind: yaml.ScalarNode,
Tag: "!!str",
Value: val.K,
}, n)
}
return &yaml.Node{
Kind: yaml.MappingNode,
Tag: "!!map",
Content: content,
}, nil
}