Return decl attrs from Value.Attribute and preserve embed YAML tags

cue/attribute.go

@@ -38,6 +38,13 @@ func (v Value) Attribute(key string) Attribute {
 		}
 		return newAttr(internal.FieldAttr, a)
 	}
+	for _, a := range export.ExtractDeclAttrs(v.v) {
+		k, _ := a.Split()
+		if key != k {
+			continue
+		}
+		return newAttr(internal.DeclAttr, a)
+	}
 
 	return nonExistAttr(key)
 }

cue/attribute_test.go

@@ -33,6 +33,14 @@ func TestAttributes(t *testing.T) {
 		@embed(foo)
 		3
 	} @field(foo)
+	d: {
+		@foo(decl=1)
+		4
+	} @foo(field=1)
+	items: [
+		"first",
+		{"second", @yaml(,tag="!Env")},
+	]
 
 	c1: {} @step(1)
 	if true {
@@ -94,6 +102,22 @@ func TestAttributes(t *testing.T) {
 		flags: cue.ValueAttr | cue.FieldAttr,
 		path:  "c4",
 		out:   "[]",
+	}, {
+		flags: cue.DeclAttr,
+		path:  "items[1]",
+		out:   `[@yaml(,tag="!Env")]`,
+	}, {
+		flags: cue.FieldAttr,
+		path:  "items[1]",
+		out:   "[]",
+	}, {
+		flags: cue.ValueAttr,
+		path:  "items[1]",
+		out:   `[@yaml(,tag="!Env")]`,
+	}, {
+		flags: cue.ValueAttr,
+		path:  "d",
+		out:   `[@foo(field=1) @foo(decl=1)]`,
 	}}
 	for _, tc := range testCases {
 		cuetdtest.FullMatrix.Run(t, tc.path, func(t *testing.T, m *cuetdtest.M) {
@@ -107,6 +131,71 @@ func TestAttributes(t *testing.T) {
 	}
 }
 
+func TestAttributeListElementDeclAttr(t *testing.T) {
+	const config = `
+	items: [
+		"first",
+		{"second", @yaml(,tag="!Env")},
+	]
+	`
+	cuetdtest.FullMatrix.Do(t, func(t *testing.T, m *cuetdtest.M) {
+		v := getValue(m, config).LookupPath(cue.ParsePath("items[1]"))
+		a := v.Attribute("yaml")
+		if err := a.Err(); err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if got, want := a.Kind(), cue.DeclAttr; got != want {
+			t.Errorf("got %v; want %v", got, want)
+		}
+		val, found, err := a.Lookup(1, "tag")
+		if err != nil {
+			t.Fatalf("lookup failed: %v", err)
+		}
+		if !found {
+			t.Fatal("tag not found")
+		}
+		if got, want := val, "!Env"; got != want {
+			t.Errorf("got %q; want %q", got, want)
+		}
+	})
+}
+
+func TestAttributePrefersFieldAttr(t *testing.T) {
+	const config = `
+	v: {
+		@foo(decl=1)
+		3
+	} @foo(field=1)
+	`
+	cuetdtest.FullMatrix.Do(t, func(t *testing.T, m *cuetdtest.M) {
+		v := getValue(m, config).Lookup("v")
+		a := v.Attribute("foo")
+		if err := a.Err(); err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if got, want := a.Kind(), cue.FieldAttr; got != want {
+			t.Errorf("got %v; want %v", got, want)
+		}
+		val, found, err := a.Lookup(0, "field")
+		if err != nil {
+			t.Fatalf("lookup failed: %v", err)
+		}
+		if !found {
+			t.Fatal("field argument not found")
+		}
+		if got, want := val, "1"; got != want {
+			t.Errorf("got %q; want %q", got, want)
+		}
+		_, found, err = a.Lookup(0, "decl")
+		if err != nil {
+			t.Fatalf("lookup failed: %v", err)
+		}
+		if found {
+			t.Fatal("unexpected decl argument")
+		}
+	})
+}
+
 func TestAttributeErr(t *testing.T) {
 	const config = `
 	a: {

internal/encoding/yaml/encode.go

@@ -234,20 +234,44 @@ func encodeExprs(exprs []ast.Expr) (n *yaml.Node, err error) {
 	return n, nil
 }
 
+// encodeEmbed encodes an embedded expression and applies any @yaml(,tag)
+// attributes found among the sibling declarations.
+func encodeEmbed(x ast.Expr, decls []ast.Decl) (*yaml.Node, error) {
+	e, err := encode(x)
+	if err != nil {
+		return nil, err
+	}
+	yamlTag, err := extractYAMLTagFromDecls(decls)
+	if err != nil {
+		return nil, err
+	}
+	if yamlTag != "" {
+		e.Tag = yamlTag
+	}
+	return e, nil
+}
+
 // extractYAMLTag looks for @yaml(,tag="...") attribute and returns the tag value.
 // Returns an empty string if no @yaml attribute or no tag argument is found.
 // Returns an error if the attribute is malformed.
 func extractYAMLTag(attrs []*ast.Attribute) (string, error) {
 	for _, attr := range attrs {
-		key, body := attr.Split()
-		if key != "yaml" {
-			continue
+		if val, found, err := yamlTagAttr(attr); err != nil {
+			return "", err
+		} else if found {
+			return val, nil
 		}
-		parsed := internal.ParseAttrBody(attr.Pos(), body)
-		if parsed.Err != nil {
-			return "", parsed.Err
+	}
+	return "", nil
+}
+
+func extractYAMLTagFromDecls(decls []ast.Decl) (string, error) {
+	for _, d := range decls {
+		attr, ok := d.(*ast.Attribute)
+		if !ok {
+			continue
 		}
-		if val, found, err := parsed.Lookup(1, "tag"); err != nil {
+		if val, found, err := yamlTagAttr(attr); err != nil {
 			return "", err
 		} else if found {
 			return val, nil
@@ -256,6 +280,26 @@ func extractYAMLTag(attrs []*ast.Attribute) (string, error) {
 	return "", nil
 }
 
+// yamlTagAttr extracts the `tag` argument from a @yaml attribute.
+func yamlTagAttr(attr *ast.Attribute) (string, bool, error) {
+	key, body := attr.Split()
+	if key != "yaml" {
+		return "", false, nil
+	}
+	parsed := internal.ParseAttrBody(attr.Pos(), body)
+	if parsed.Err != nil {
+		return "", false, parsed.Err
+	}
+	val, found, err := parsed.Lookup(1, "tag")
+	if err != nil {
+		return "", false, err
+	}
+	if !found {
+		return "", false, nil
+	}
+	return val, true, nil
+}
+
 // encodeDecls converts a sequence of declarations to a value. If it encounters
 // an embedded value, it will return this expression. This is more relaxed for
 // structs than is currently allowed for CUE, but the expectation is that this
@@ -331,14 +375,29 @@ func encodeDecls(decls []ast.Decl) (n *yaml.Node, err error) {
 				return nil, errors.Newf(x.Pos(), "yaml: multiple embedded values")
 			}
 			hasEmbed = true
-			e, err := encode(x.Expr)
+			e, err := encodeEmbed(x.Expr, decls)
 			if err != nil {
 				return nil, err
 			}
 			addDocs(x, e, e)
 			lastHead = e
 			lastFoot = e
 			n.Content = append(n.Content, e)
+
+		case ast.Expr:
+			// cue.Value.Syntax may produce bare expressions rather
+			// than EmbedDecl wrappers in struct element lists.
+			if hasEmbed {
+				return nil, errors.Newf(d.Pos(), "yaml: multiple embedded values")
+			}
+			hasEmbed = true
+			e, err := encodeEmbed(x, decls)
+			if err != nil {
+				return nil, err
+			}
+			lastHead = e
+			lastFoot = e
+			n.Content = append(n.Content, e)
 		}
 		if docForNext.Len() > 0 {
 			docForNext.WriteString(lastHead.HeadComment)

internal/encoding/yaml/encode_test.go

@@ -153,6 +153,14 @@ f4: {} # line 4
 # hex
 0xabc # line
 # trail
+`,
+	}, {
+		name: "embed_binary",
+		in: `
+		{'\x80'}
+		`,
+		out: `
+!!binary gA==
 `,
 	}, {
 		// TODO: support this at some point
@@ -307,6 +315,19 @@ custom: !%3Ctag:example.com,2000:app/foo%3E value
 		out: `
 item: !%3Chttps://example.com/schema/v1%3E value
 		`,
+	}, {
+		name: "yaml_tag_embed",
+		in: `
+		items: [
+			"first",
+			{"second", @yaml(,tag="!Env")},
+		] @yaml(,tag="!Format")
+		`,
+		out: `
+items: !Format
+  - first
+  - !Env second
+		`,
 	}, {
 		name: "yaml_attribute_without_tag",
 		in: `