yq/pkg/yqlib/encoder.go

248 lines
5.1 KiB
Go
Raw Normal View History

2020-01-10 11:01:59 +00:00
package yqlib
import (
2020-02-27 23:09:49 +00:00
"bytes"
2020-01-10 11:01:59 +00:00
"encoding/json"
"fmt"
2020-01-10 11:01:59 +00:00
"io"
yaml "gopkg.in/yaml.v3"
)
type Encoder interface {
Encode(node *yaml.Node) error
}
type yamlEncoder struct {
2020-02-27 23:09:49 +00:00
destination io.Writer
indent int
colorise bool
firstDoc bool
2020-01-10 11:01:59 +00:00
}
2020-02-27 23:09:49 +00:00
func NewYamlEncoder(destination io.Writer, indent int, colorise bool) Encoder {
2020-02-03 05:52:12 +00:00
if indent < 0 {
indent = 0
}
2020-02-27 23:09:49 +00:00
return &yamlEncoder{destination, indent, colorise, true}
2020-01-10 11:01:59 +00:00
}
func (ye *yamlEncoder) Encode(node *yaml.Node) error {
2020-02-27 23:09:49 +00:00
destination := ye.destination
tempBuffer := bytes.NewBuffer(nil)
if ye.colorise {
destination = tempBuffer
}
var encoder = yaml.NewEncoder(destination)
encoder.SetIndent(ye.indent)
// TODO: work out if the first doc had a separator or not.
if ye.firstDoc {
ye.firstDoc = false
} else if _, err := destination.Write([]byte("---\n")); err != nil {
return err
}
if err := encoder.Encode(node); err != nil {
return err
}
if ye.colorise {
return ColorizeAndPrint(tempBuffer.Bytes(), ye.destination)
}
return nil
2020-01-10 11:01:59 +00:00
}
type jsonEncoder struct {
encoder *json.Encoder
}
func mapKeysToStrings(node *yaml.Node) {
2020-09-13 00:59:40 +00:00
if node.Kind == yaml.MappingNode {
for index, child := range node.Content {
2020-09-13 00:59:40 +00:00
if index%2 == 0 { // its a map key
child.Tag = "!!str"
}
2020-09-13 00:59:40 +00:00
}
}
for _, child := range node.Content {
mapKeysToStrings(child)
}
}
2020-11-03 23:48:43 +00:00
func NewJsonEncoder(destination io.Writer, indent int) Encoder {
2020-01-10 11:01:59 +00:00
var encoder = json.NewEncoder(destination)
2020-02-03 05:52:12 +00:00
var indentString = ""
for index := 0; index < indent; index++ {
indentString = indentString + " "
}
2020-11-03 23:48:43 +00:00
encoder.SetIndent("", indentString)
2020-01-10 11:01:59 +00:00
return &jsonEncoder{encoder}
}
func (je *jsonEncoder) Encode(node *yaml.Node) error {
var dataBucket orderedMap
// firstly, convert all map keys to strings
mapKeysToStrings(node)
2020-01-10 11:01:59 +00:00
errorDecoding := node.Decode(&dataBucket)
if errorDecoding != nil {
return errorDecoding
}
return je.encoder.Encode(dataBucket)
}
// 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(); 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 && 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) {
if o.kv == nil {
return json.Marshal(o.altVal)
}
buf := new(bytes.Buffer)
enc := json.NewEncoder(buf)
buf.WriteByte('{')
for idx, el := range o.kv {
if err := enc.Encode(el.K); err != nil {
return nil, err
}
buf.WriteByte(':')
2020-10-07 17:22:45 +00:00
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:
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
}