Merge pull request #59 from goccy/feature/fix-indent-for-encoding

Fix calculating indent number at encoding

ast/ast.go

@@ -102,6 +102,8 @@ type Node interface {
 	GetToken() *token.Token
 	// Type returns type of node
 	Type() NodeType
+	// AddColumn add column number to child nodes recursively
+	AddColumn(int)
 }
 
 // File contains all documents in YAML file
@@ -131,6 +133,13 @@ func (d *Document) GetToken() *token.Token {
 	return d.Body.GetToken()
 }
 
+// AddColumn add column number to child nodes recursively
+func (d *Document) AddColumn(col int) {
+	if d.Body != nil {
+		d.Body.AddColumn(col)
+	}
+}
+
 // Type returns DocumentType
 func (d *Document) Type() NodeType { return DocumentType }
 
@@ -310,6 +319,11 @@ func (n *NullNode) GetToken() *token.Token {
 	return n.Token
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *NullNode) AddColumn(col int) {
+	n.Token.AddColumn(col)
+}
+
 // GetValue returns nil value
 func (n *NullNode) GetValue() interface{} {
 	return nil
@@ -335,6 +349,11 @@ func (n *IntegerNode) GetToken() *token.Token {
 	return n.Token
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *IntegerNode) AddColumn(col int) {
+	n.Token.AddColumn(col)
+}
+
 // GetValue returns int64 value
 func (n *IntegerNode) GetValue() interface{} {
 	return n.Value
@@ -361,6 +380,11 @@ func (n *FloatNode) GetToken() *token.Token {
 	return n.Token
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *FloatNode) AddColumn(col int) {
+	n.Token.AddColumn(col)
+}
+
 // GetValue returns float64 value
 func (n *FloatNode) GetValue() interface{} {
 	return n.Value
@@ -386,6 +410,11 @@ func (n *StringNode) GetToken() *token.Token {
 	return n.Token
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *StringNode) AddColumn(col int) {
+	n.Token.AddColumn(col)
+}
+
 // GetValue returns string value
 func (n *StringNode) GetValue() interface{} {
 	return n.Value
@@ -417,6 +446,14 @@ func (n *LiteralNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *LiteralNode) AddColumn(col int) {
+	n.Start.AddColumn(col)
+	if n.Value != nil {
+		n.Value.AddColumn(col)
+	}
+}
+
 // GetValue returns string value
 func (n *LiteralNode) GetValue() interface{} {
 	return n.Value.GetValue()
@@ -467,6 +504,11 @@ func (n *BoolNode) GetToken() *token.Token {
 	return n.Token
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *BoolNode) AddColumn(col int) {
+	n.Token.AddColumn(col)
+}
+
 // GetValue returns boolean value
 func (n *BoolNode) GetValue() interface{} {
 	return n.Value
@@ -492,6 +534,11 @@ func (n *InfinityNode) GetToken() *token.Token {
 	return n.Token
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *InfinityNode) AddColumn(col int) {
+	n.Token.AddColumn(col)
+}
+
 // GetValue returns math.Inf(0) or math.Inf(-1)
 func (n *InfinityNode) GetValue() interface{} {
 	return n.Value
@@ -516,6 +563,11 @@ func (n *NanNode) GetToken() *token.Token {
 	return n.Token
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *NanNode) AddColumn(col int) {
+	n.Token.AddColumn(col)
+}
+
 // GetValue returns math.NaN()
 func (n *NanNode) GetValue() interface{} {
 	return math.NaN()
@@ -575,6 +627,15 @@ func (n *MappingNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *MappingNode) AddColumn(col int) {
+	n.Start.AddColumn(col)
+	n.End.AddColumn(col)
+	for _, value := range n.Values {
+		value.AddColumn(col)
+	}
+}
+
 func (n *MappingNode) flowStyleString() string {
 	if len(n.Values) == 0 {
 		return "{}"
@@ -628,6 +689,17 @@ func (n *MappingValueNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *MappingValueNode) AddColumn(col int) {
+	n.Start.AddColumn(col)
+	if n.Key != nil {
+		n.Key.AddColumn(col)
+	}
+	if n.Value != nil {
+		n.Value.AddColumn(col)
+	}
+}
+
 // String mapping value to text
 func (n *MappingValueNode) String() string {
 	space := strings.Repeat(" ", n.Key.GetToken().Position.Column-1)
@@ -702,6 +774,15 @@ func (n *SequenceNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *SequenceNode) AddColumn(col int) {
+	n.Start.AddColumn(col)
+	n.End.AddColumn(col)
+	for _, value := range n.Values {
+		value.AddColumn(col)
+	}
+}
+
 func (n *SequenceNode) flowStyleString() string {
 	values := []string{}
 	for _, value := range n.Values {
@@ -760,6 +841,17 @@ func (n *AnchorNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *AnchorNode) AddColumn(col int) {
+	n.Start.AddColumn(col)
+	if n.Name != nil {
+		n.Name.AddColumn(col)
+	}
+	if n.Value != nil {
+		n.Value.AddColumn(col)
+	}
+}
+
 // String anchor to text
 func (n *AnchorNode) String() string {
 	value := n.Value.String()
@@ -783,6 +875,14 @@ func (n *AliasNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *AliasNode) AddColumn(col int) {
+	n.Start.AddColumn(col)
+	if n.Value != nil {
+		n.Value.AddColumn(col)
+	}
+}
+
 // String alias to text
 func (n *AliasNode) String() string {
 	return fmt.Sprintf("*%s", n.Value.String())
@@ -802,6 +902,13 @@ func (n *DirectiveNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *DirectiveNode) AddColumn(col int) {
+	if n.Value != nil {
+		n.Value.AddColumn(col)
+	}
+}
+
 // String directive to text
 func (n *DirectiveNode) String() string {
 	return fmt.Sprintf("%s%s", n.Start.Value, n.Value.String())
@@ -821,6 +928,14 @@ func (n *TagNode) GetToken() *token.Token {
 	return n.Start
 }
 
+// AddColumn add column number to child nodes recursively
+func (n *TagNode) AddColumn(col int) {
+	n.Start.AddColumn(col)
+	if n.Value != nil {
+		n.Value.AddColumn(col)
+	}
+}
+
 // String tag to text
 func (n *TagNode) String() string {
 	return fmt.Sprintf("%s %s", n.Start.Value, n.Value.String())

encode.go

@@ -249,9 +249,7 @@ func (e *Encoder) encodeMapItem(item MapItem, column int) (*ast.MappingValueNode
 		return nil, errors.Wrapf(err, "failed to encode MapItem")
 	}
 	if m, ok := value.(*ast.MappingNode); ok {
-		for _, value := range m.Values {
-			value.Key.GetToken().Position.Column += e.indent
-		}
+		m.AddColumn(e.indent)
 	}
 	return &ast.MappingValueNode{
 		Start: token.New("", "", e.pos(column)),
@@ -287,9 +285,7 @@ func (e *Encoder) encodeMap(value reflect.Value, column int) ast.Node {
 			return nil
 		}
 		if m, ok := value.(*ast.MappingNode); ok {
-			for _, value := range m.Values {
-				value.Key.GetToken().Position.Column += e.indent
-			}
+			m.AddColumn(e.indent)
 		}
 		node.Values = append(node.Values, &ast.MappingValueNode{
 			Key:   e.encodeString(k.Interface().(string), column),
@@ -377,10 +373,7 @@ func (e *Encoder) encodeStruct(value reflect.Value, column int) (ast.Node, error
 			if !e.isFlowStyle && structField.IsFlow {
 				m.IsFlowStyle = true
 			}
-			for _, value := range m.Values {
-				value.Key.GetToken().Position.Column += e.indent
-				value.Value.GetToken().Position.Column += e.indent
-			}
+			value.AddColumn(e.indent)
 		} else if s, ok := value.(*ast.SequenceNode); ok {
 			if !e.isFlowStyle && structField.IsFlow {
 				s.IsFlowStyle = true
@@ -454,8 +447,8 @@ func (e *Encoder) encodeStruct(value reflect.Value, column int) (ast.Node, error
 					// if declared same key name, skip encoding this field
 					continue
 				}
-				key.GetToken().Position.Column -= e.indent
-				value.GetToken().Position.Column -= e.indent
+				key.AddColumn(-e.indent)
+				value.AddColumn(-e.indent)
 				node.Values = append(node.Values, &ast.MappingValueNode{
 					Key:   key,
 					Value: value,

encode_test.go

@@ -319,6 +319,28 @@ func TestEncoder(t *testing.T) {
 	}
 }
 
+func TestEncodeStructIncludeMap(t *testing.T) {
+	type U struct {
+		M map[string]string
+	}
+	type T struct {
+		A U
+	}
+	bytes, err := yaml.Marshal(T{
+		A: U{
+			M: map[string]string{"x": "y"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("%+v", err)
+	}
+	expect := "a:\n  m:\n    x: y\n"
+	actual := string(bytes)
+	if actual != expect {
+		t.Fatalf("unexpected output. expect:[%s] actual:[%s]", expect, actual)
+	}
+}
+
 func TestEncodeWithAnchorAndAlias(t *testing.T) {
 	var buf bytes.Buffer
 	enc := yaml.NewEncoder(&buf)

token/token.go

@@ -664,6 +664,14 @@ func (t *Token) NextType() Type {
 	return UnknownType
 }
 
+// AddColumn append column number to current position of column
+func (t *Token) AddColumn(col int) {
+	if t == nil {
+		return
+	}
+	t.Position.Column += col
+}
+
 // Tokens type of token collection
 type Tokens []*Token
 

yaml_test.go

@@ -52,6 +52,9 @@ func (t *marshalTest) MarshalYAML() ([]byte, error) {
 		{
 			"c", true,
 		},
+		{
+			"d", map[string]string{"x": "y"},
+		},
 	})
 }
 
@@ -87,6 +90,8 @@ a:
   a: 1
   b: hello
   c: true
+  d:
+    x: y
 b:
   a: 2
   b: world