/
yaml_encoder.go
99 lines (84 loc) · 2.77 KB
/
yaml_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
package serde
import (
"bytes"
"encoding/json"
"fmt"
"io"
"github.com/gogo/protobuf/jsonpb"
"github.com/gogo/protobuf/proto"
"gopkg.in/pachyderm/yaml.v3"
)
// YAMLEncoder is an implementation of serde.Encoder that operates on YAML data
type YAMLEncoder struct {
e *yaml.Encoder
// OrigName sets whether this YAMLEncoder uses the original (.proto) name of
// fields when marshalling to protos
origName bool
}
// EncodeYAML is a convenience function that encodes yaml data using a
// YAMLEncoder, but can be called inline
func EncodeYAML(v interface{}, options ...EncoderOption) ([]byte, error) {
var buf bytes.Buffer
e := NewYAMLEncoder(&buf, options...)
if err := e.Encode(v); err != nil {
return nil, err
}
return buf.Bytes(), nil
}
// NewYAMLEncoder returns a new YAMLEncoder that writes to 'w'
func NewYAMLEncoder(w io.Writer, options ...EncoderOption) *YAMLEncoder {
e := &YAMLEncoder{e: yaml.NewEncoder(w)}
for _, o := range options {
o(e)
}
return e
}
// Encode implements the corresponding method of serde.Encoder
func (e *YAMLEncoder) Encode(v interface{}) error {
return e.EncodeTransform(v, nil)
}
// EncodeTransform implements the corresponding method of serde.Encoder
func (e *YAMLEncoder) EncodeTransform(v interface{}, f func(map[string]interface{}) error) error {
// Encode to JSON first
var buf bytes.Buffer
j := json.NewEncoder(&buf)
if err := j.Encode(v); err != nil {
return fmt.Errorf("serialization error while canonicalizing output: %v", err)
}
return e.jsonToYAMLTransform(buf.Bytes(), f)
}
// EncodeProto implements the corresponding method of serde.Encoder
func (e *YAMLEncoder) EncodeProto(v proto.Message) error {
return e.EncodeProtoTransform(v, nil)
}
// EncodeProtoTransform implements the corresponding method of serde.Encoder
func (e *YAMLEncoder) EncodeProtoTransform(v proto.Message, f func(map[string]interface{}) error) error {
// Encode to JSON first
var buf bytes.Buffer
m := jsonpb.Marshaler{
OrigName: e.origName,
}
if err := m.Marshal(&buf, v); err != nil {
return fmt.Errorf("serialization error while canonicalizing output: %v", err)
}
return e.jsonToYAMLTransform(buf.Bytes(), f)
}
func (e *YAMLEncoder) jsonToYAMLTransform(intermediateJSON []byte,
f func(map[string]interface{}) error) error {
// Unmarshal from JSON to intermediate map ('holder')
holder := map[string]interface{}{}
if err := json.Unmarshal(intermediateJSON, &holder); err != nil {
return fmt.Errorf("deserialization error while canonicalizing output: %v", err)
}
// transform 'holder' (e.g. de-stringifying TFJob)
if f != nil {
if err := f(holder); err != nil {
return err
}
}
// Encode 'holder' to YAML
if err := e.e.Encode(holder); err != nil {
return fmt.Errorf("serialization error while canonicalizing yaml: %v", err)
}
return nil
}