Skip to content

Commit

Permalink
Support MarshalAnchor option for encoder
Browse files Browse the repository at this point in the history
  • Loading branch information
goccy committed Jan 9, 2020
1 parent 009880f commit 3544475
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 17 deletions.
45 changes: 29 additions & 16 deletions encode.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ type Encoder struct {
opts []EncodeOption
indent int
isFlowStyle bool
anchorCallback func(*ast.AnchorNode, interface{}) error
anchorPtrToNameMap map[uintptr]string

line int
Expand Down Expand Up @@ -347,6 +348,26 @@ func (e *Encoder) encodeTime(v time.Time, column int) ast.Node {
return ast.String(token.New(value, value, e.pos(column)))
}

func (e *Encoder) encodeAnchor(anchorName string, value ast.Node, fieldValue reflect.Value, column int) (ast.Node, error) {
anchorNode := &ast.AnchorNode{
Start: token.New("&", "&", e.pos(column)),
Name: ast.String(token.New(anchorName, anchorName, e.pos(column))),
Value: value,
}
if e.anchorCallback != nil {
if err := e.anchorCallback(anchorNode, fieldValue.Interface()); err != nil {
return nil, errors.Wrapf(err, "failed to marshal anchor")
}
if snode, ok := anchorNode.Name.(*ast.StringNode); ok {
anchorName = snode.Value
}
}
if fieldValue.Kind() == reflect.Ptr {
e.anchorPtrToNameMap[fieldValue.Pointer()] = anchorName
}
return anchorNode, nil
}

func (e *Encoder) encodeStruct(value reflect.Value, column int) (ast.Node, error) {
node := ast.Mapping(token.New("", "", e.pos(column)), e.isFlowStyle)
structType := value.Type()
Expand Down Expand Up @@ -382,25 +403,17 @@ func (e *Encoder) encodeStruct(value reflect.Value, column int) (ast.Node, error
key := e.encodeString(structField.RenderName, column)
switch {
case structField.AnchorName != "":
anchorName := structField.AnchorName
if fieldValue.Kind() == reflect.Ptr {
e.anchorPtrToNameMap[fieldValue.Pointer()] = anchorName
}
value = &ast.AnchorNode{
Start: token.New("&", "&", e.pos(column)),
Name: ast.String(token.New(anchorName, anchorName, e.pos(column))),
Value: value,
anchorNode, err := e.encodeAnchor(structField.AnchorName, value, fieldValue, column)
if err != nil {
return nil, errors.Wrapf(err, "failed to encode anchor")
}
value = anchorNode
case structField.IsAutoAnchor:
anchorName := structField.RenderName
if fieldValue.Kind() == reflect.Ptr {
e.anchorPtrToNameMap[fieldValue.Pointer()] = anchorName
}
value = &ast.AnchorNode{
Start: token.New("&", "&", e.pos(column)),
Name: ast.String(token.New(anchorName, anchorName, e.pos(column))),
Value: value,
anchorNode, err := e.encodeAnchor(structField.RenderName, value, fieldValue, column)
if err != nil {
return nil, errors.Wrapf(err, "failed to encode anchor")
}
value = anchorNode
case structField.IsAutoAlias:
if fieldValue.Kind() != reflect.Ptr {
return nil, xerrors.Errorf(
Expand Down
80 changes: 80 additions & 0 deletions encode_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"time"

"github.com/goccy/go-yaml"
"github.com/goccy/go-yaml/ast"
)

func TestEncoder(t *testing.T) {
Expand Down Expand Up @@ -567,6 +568,85 @@ func TestEncoder_Flow(t *testing.T) {
}
}

func TestEncoder_MarshalAnchor(t *testing.T) {
type Host struct {
Hostname string
Username string
Password string
}
type HostDecl struct {
Host *Host `yaml:",anchor"`
}
type Queue struct {
Name string `yaml:","`
*Host `yaml:",alias"`
}
var doc struct {
Hosts []*HostDecl `yaml:"hosts"`
Queues []*Queue `yaml:"queues"`
}
host1 := &Host{
Hostname: "host1.example.com",
Username: "userA",
Password: "pass1",
}
host2 := &Host{
Hostname: "host2.example.com",
Username: "userB",
Password: "pass2",
}
doc.Hosts = []*HostDecl{
{
Host: host1,
},
{
Host: host2,
},
}
doc.Queues = []*Queue{
{
Name: "queue",
Host: host1,
}, {
Name: "queue2",
Host: host2,
},
}
hostIdx := 1
opt := yaml.MarshalAnchor(func(anchor *ast.AnchorNode, value interface{}) error {
if _, ok := value.(*Host); ok {
nameNode := anchor.Name.(*ast.StringNode)
nameNode.Value = fmt.Sprintf("host%d", hostIdx)
hostIdx++
}
return nil
})

var buf bytes.Buffer
if err := yaml.NewEncoder(&buf, opt).Encode(doc); err != nil {
t.Fatalf("%+v", err)
}
expect := `
hosts:
- host: &host1
hostname: host1.example.com
username: userA
password: pass1
- host: &host2
hostname: host2.example.com
username: userB
password: pass2
queues:
- name: queue
host: *host1
- name: queue2
host: *host2
`
if "\n"+buf.String() != expect {
t.Fatalf("unexpected output. %s", buf.String())
}
}

func Example_Marshal_ExplicitAnchorAlias() {
type T struct {
A int
Expand Down
14 changes: 13 additions & 1 deletion option.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
package yaml

import "io"
import (
"io"

"github.com/goccy/go-yaml/ast"
)

// DecodeOption functional option type for Decoder
type DecodeOption func(d *Decoder) error
Expand Down Expand Up @@ -73,3 +77,11 @@ func Flow(isFlowStyle bool) EncodeOption {
return nil
}
}

// MarshalAnchor call back if encoder find an anchor during encoding
func MarshalAnchor(callback func(*ast.AnchorNode, interface{}) error) EncodeOption {
return func(e *Encoder) error {
e.anchorCallback = callback
return nil
}
}

0 comments on commit 3544475

Please sign in to comment.