Improve SetString and fix !!binary auto-detection.

emitterc.go

@@ -1920,7 +1920,6 @@ func yaml_emitter_write_folded_scalar(emitter *yaml_emitter_t, value []byte) boo
 }
 
 func yaml_emitter_write_comment(emitter *yaml_emitter_t, comment []byte) bool {
-	// [Go] TODO Emit "# " when necessary.
 	breaks := false
 	pound := false
 	for i := 0; i < len(comment); {

encode.go

@@ -489,6 +489,20 @@ func (e *encoder) node(node *Node) {
 		e.emit()
 
 	case ScalarNode:
+		value := node.Value
+		if !utf8.ValidString(value) {
+			if tag == binaryTag {
+				failf("explicitly tagged !!binary data must be base64-encoded")
+			}
+			if tag != "" {
+				failf("cannot marshal invalid UTF-8 data as %s", shortTag(tag))
+			}
+			// It can't be encoded directly as YAML so use a binary tag
+			// and encode it as base64.
+			tag = binaryTag
+			value = encodeBase64(value)
+		}
+
 		style := yaml_PLAIN_SCALAR_STYLE
 		switch {
 		case node.Style&DoubleQuotedStyle != 0:
@@ -499,27 +513,12 @@ func (e *encoder) node(node *Node) {
 			style = yaml_LITERAL_SCALAR_STYLE
 		case node.Style&FoldedStyle != 0:
 			style = yaml_FOLDED_SCALAR_STYLE
-		case strings.Contains(node.Value, "\n"):
+		case strings.Contains(value, "\n"):
 			style = yaml_LITERAL_SCALAR_STYLE
 		case forceQuoting:
 			style = yaml_DOUBLE_QUOTED_SCALAR_STYLE
 		}
 
-		e.emitScalar(node.Value, node.Anchor, tag, style, []byte(node.HeadComment), []byte(node.LineComment), []byte(node.FootComment))
-
-		// TODO Check if binaries are being decoded into node.Value or not.
-		//switch {
-		//if !utf8.ValidString(s) {
-		//	if tag == binaryTag {
-		//		failf("explicitly tagged !!binary data must be base64-encoded")
-		//	}
-		//	if tag != "" {
-		//		failf("cannot marshal invalid UTF-8 data as %s", shortTag(tag))
-		//	}
-		//	// It can't be encoded directly as YAML so use a binary tag
-		//	// and encode it as base64.
-		//	tag = binaryTag
-		//	s = encodeBase64(s)
-		//}
+		e.emitScalar(value, node.Anchor, tag, style, []byte(node.HeadComment), []byte(node.LineComment), []byte(node.FootComment))
 	}
 }

node_test.go

@@ -116,6 +116,20 @@ var nodeTests = []struct {
 				Column: 1,
 			}},
 		},
+	}, {
+		// Item doesn't have a tag, but needs to be binary encoded due to its content.
+		"[encode]!!binary gIGC\n",
+		yaml.Node{
+			Kind:   yaml.DocumentNode,
+			Line:   1,
+			Column: 1,
+			Content: []*yaml.Node{{
+				Kind:   yaml.ScalarNode,
+				Value:  "\x80\x81\x82",
+				Line:   1,
+				Column: 1,
+			}},
+		},
 	}, {
 		// Same, but with strings we can just quote them.
 		"[encode]\"123\"\n",
@@ -1157,3 +1171,93 @@ func dropNode(name string) *yaml.Node {
 	delete(savedNodes, name)
 	return node
 }
+
+var setStringTests = []struct {
+	str  string
+	yaml string
+	node yaml.Node
+}{
+	{
+		"something simple",
+		"something simple\n",
+		yaml.Node{
+			Kind:   yaml.ScalarNode,
+			Value:  "something simple",
+			Tag:    "!!str",
+		},
+	}, {
+		`"quoted value"`,
+		"'\"quoted value\"'\n",
+		yaml.Node{
+			Kind:   yaml.ScalarNode,
+			Value:  `"quoted value"`,
+			Tag:    "!!str",
+		},
+	}, {
+		"multi\nline",
+		"|-\n  multi\n  line\n",
+		yaml.Node{
+			Kind:   yaml.ScalarNode,
+			Value:  "multi\nline",
+			Tag:    "!!str",
+			Style:  yaml.LiteralStyle,
+		},
+	}, {
+		"123",
+		"\"123\"\n",
+		yaml.Node{
+			Kind:   yaml.ScalarNode,
+			Value:  "123",
+			Tag:    "!!str",
+		},
+	}, {
+		"multi\nline\n",
+		"|\n  multi\n  line\n",
+		yaml.Node{
+			Kind:   yaml.ScalarNode,
+			Value:  "multi\nline\n",
+			Tag:    "!!str",
+			Style:  yaml.LiteralStyle,
+		},
+	}, {
+		"\x80\x81\x82",
+		"!!binary gIGC\n",
+		yaml.Node{
+			Kind:   yaml.ScalarNode,
+			Value:  "gIGC",
+			Tag:    "!!binary",
+		},
+	},
+}
+
+func (s *S) TestSetString(c *C) {
+	defer os.Setenv("TZ", os.Getenv("TZ"))
+	os.Setenv("TZ", "UTC")
+	for i, item := range setStringTests {
+		c.Logf("test %d: %q", i, item.str)
+
+		var node yaml.Node
+
+		node.SetString(item.str)
+
+		c.Assert(node, DeepEquals, item.node)
+
+		buf := bytes.Buffer{}
+		enc := yaml.NewEncoder(&buf)
+		enc.SetIndent(2)
+		err := enc.Encode(&item.node)
+		c.Assert(err, IsNil)
+		err = enc.Close()
+		c.Assert(err, IsNil)
+		c.Assert(buf.String(), Equals, item.yaml)
+
+		var doc yaml.Node
+		err = yaml.Unmarshal([]byte(item.yaml), &doc)
+		c.Assert(err, IsNil)
+
+		var str string
+		err = node.Decode(&str)
+		c.Assert(err, IsNil)
+		c.Assert(str, Equals, item.str)
+	}
+}

yaml.go

@@ -28,6 +28,7 @@ import (
 	"reflect"
 	"strings"
 	"sync"
+	"unicode/utf8"
 )
 
 // The Unmarshaler interface may be implemented by types to customize their
@@ -417,11 +418,15 @@ func (n *Node) indicatedString() bool {
 // and defines its style in a pleasant way depending on its content.
 func (n *Node) SetString(s string) {
 	n.Kind = ScalarNode
-	n.Value = s
-	if strings.Contains(s, "\n") {
+	if utf8.ValidString(s) {
+		n.Value = s
+		n.Tag = strTag
+	} else {
+		n.Value = encodeBase64(s)
+		n.Tag = binaryTag
+	}
+	if strings.Contains(n.Value, "\n") {
 		n.Style = LiteralStyle
-	} else if n.ShortTag() != strTag {
-		n.Style = DoubleQuotedStyle
 	}
 }