hclwrite: Implement Body.SetAttributeTraversal

hclwrite/ast_body.go

@@ -87,7 +87,16 @@ func (b *Body) SetAttributeValue(name string, val cty.Value) *Attribute {
 // The return value is the attribute that was either modified in-place or
 // created.
 func (b *Body) SetAttributeTraversal(name string, traversal hcl.Traversal) *Attribute {
-	panic("Body.SetAttributeTraversal not yet implemented")
+	attr := b.GetAttribute(name)
+	expr := NewExpressionAbsTraversal(traversal)
+	if attr != nil {
+		attr.expr = attr.expr.ReplaceWith(expr)
+	} else {
+		attr := newAttribute()
+		attr.init(name, expr)
+		b.appendItem(attr)
+	}
+	return attr
 }
 
 // AppendBlock appends an existing block (which must not be already attached

hclwrite/ast_body_test.go

@@ -414,6 +414,232 @@ func TestBodySetAttributeValue(t *testing.T) {
 	}
 }
 
+func TestBodySetAttributeTraversal(t *testing.T) {
+	tests := []struct {
+		src  string
+		name string
+		trav string
+		want Tokens
+	}{
+		{
+			"",
+			"a",
+			`b`,
+			Tokens{
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte{'a'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEqual,
+					Bytes:        []byte{'='},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte("b"),
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenNewline,
+					Bytes:        []byte{'\n'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+		{
+			"",
+			"a",
+			`b.c.d`,
+			Tokens{
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte{'a'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEqual,
+					Bytes:        []byte{'='},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte("b"),
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenDot,
+					Bytes:        []byte("."),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte("c"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenDot,
+					Bytes:        []byte("."),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte("d"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenNewline,
+					Bytes:        []byte{'\n'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+		{
+			"",
+			"a",
+			`b[0]`,
+			Tokens{
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte{'a'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEqual,
+					Bytes:        []byte{'='},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte("b"),
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenOBrack,
+					Bytes:        []byte("["),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenNumberLit,
+					Bytes:        []byte("0"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenCBrack,
+					Bytes:        []byte("]"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenNewline,
+					Bytes:        []byte{'\n'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+		{
+			"",
+			"a",
+			`b[0].c`,
+			Tokens{
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte{'a'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEqual,
+					Bytes:        []byte{'='},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte("b"),
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenOBrack,
+					Bytes:        []byte("["),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenNumberLit,
+					Bytes:        []byte("0"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenCBrack,
+					Bytes:        []byte("]"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenDot,
+					Bytes:        []byte("."),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte("c"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenNewline,
+					Bytes:        []byte{'\n'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%s = %s in %s", test.name, test.trav, test.src), func(t *testing.T) {
+			f, diags := ParseConfig([]byte(test.src), "", hcl.Pos{Line: 1, Column: 1})
+			if len(diags) != 0 {
+				for _, diag := range diags {
+					t.Logf("- %s", diag.Error())
+				}
+				t.Fatalf("unexpected diagnostics")
+			}
+
+			traversal, diags := hclsyntax.ParseTraversalAbs([]byte(test.trav), "", hcl.Pos{Line: 1, Column: 1})
+			if len(diags) != 0 {
+				for _, diag := range diags {
+					t.Logf("- %s", diag.Error())
+				}
+				t.Fatalf("unexpected diagnostics from traversal")
+			}
+
+			f.Body().SetAttributeTraversal(test.name, traversal)
+			got := f.BuildTokens(nil)
+			format(got)
+			if !reflect.DeepEqual(got, test.want) {
+				diff := cmp.Diff(test.want, got)
+				t.Errorf("wrong result\ngot:  %s\nwant: %s\ndiff:\n%s", spew.Sdump(got), spew.Sdump(test.want), diff)
+			}
+		})
+	}
+}
+
 func TestBodyAppendBlock(t *testing.T) {
 	tests := []struct {
 		src       string

hclwrite/ast_expression.go

@@ -2,6 +2,7 @@ package hclwrite
 
 import (
 	"github.com/hashicorp/hcl2/hcl"
+	"github.com/hashicorp/hcl2/hcl/hclsyntax"
 	"github.com/zclconf/go-cty/cty"
 )
 
@@ -41,7 +42,61 @@ func NewExpressionLiteral(val cty.Value) *Expression {
 // NewExpressionAbsTraversal constructs an expression that represents the
 // given traversal, which must be absolute or this function will panic.
 func NewExpressionAbsTraversal(traversal hcl.Traversal) *Expression {
-	panic("NewExpressionAbsTraversal not yet implemented")
+	if traversal.IsRelative() {
+		panic("can't construct expression from relative traversal")
+	}
+
+	physT := newTraversal()
+	rootName := traversal.RootName()
+	steps := traversal[1:]
+
+	{
+		tn := newTraverseName()
+		tn.name = tn.children.Append(newIdentifier(&Token{
+			Type:  hclsyntax.TokenIdent,
+			Bytes: []byte(rootName),
+		}))
+		physT.steps.Add(physT.children.Append(tn))
+	}
+
+	for _, step := range steps {
+		switch ts := step.(type) {
+		case hcl.TraverseAttr:
+			tn := newTraverseName()
+			tn.children.AppendUnstructuredTokens(Tokens{
+				{
+					Type:  hclsyntax.TokenDot,
+					Bytes: []byte{'.'},
+				},
+			})
+			tn.name = tn.children.Append(newIdentifier(&Token{
+				Type:  hclsyntax.TokenIdent,
+				Bytes: []byte(ts.Name),
+			}))
+			physT.steps.Add(physT.children.Append(tn))
+		case hcl.TraverseIndex:
+			ti := newTraverseIndex()
+			ti.children.AppendUnstructuredTokens(Tokens{
+				{
+					Type:  hclsyntax.TokenOBrack,
+					Bytes: []byte{'['},
+				},
+			})
+			indexExpr := NewExpressionLiteral(ts.Key)
+			ti.key = ti.children.Append(indexExpr)
+			ti.children.AppendUnstructuredTokens(Tokens{
+				{
+					Type:  hclsyntax.TokenCBrack,
+					Bytes: []byte{']'},
+				},
+			})
+			physT.steps.Add(physT.children.Append(ti))
+		}
+	}
+
+	expr := newExpression()
+	expr.absTraversals.Add(expr.children.Append(physT))
+	return expr
 }
 
 type Traversal struct {

hclwrite/examples_test.go

@@ -3,6 +3,7 @@ package hclwrite_test
 import (
 	"fmt"
 
+	"github.com/hashicorp/hcl2/hcl"
 	"github.com/hashicorp/hcl2/hclwrite"
 	"github.com/zclconf/go-cty/cty"
 )
@@ -19,6 +20,14 @@ func Example_generateFromScratch() {
 	}))
 	rootBody.SetAttributeValue("string", cty.StringVal("foo"))
 	rootBody.SetAttributeValue("bool", cty.False)
+	rootBody.SetAttributeTraversal("path", hcl.Traversal{
+		hcl.TraverseRoot{
+			Name: "env",
+		},
+		hcl.TraverseAttr{
+			Name: "PATH",
+		},
+	})
 	rootBody.AppendNewline()
 	fooBlock := rootBody.AppendNewBlock("foo", nil)
 	fooBody := fooBlock.Body()
@@ -41,6 +50,7 @@ func Example_generateFromScratch() {
 	//
 	// object = {bar = 5, baz = true, foo = "foo"}
 	// bool   = false
+	// path   = env.PATH
 	//
 	// foo {
 	//   hello = "world"

hclwrite/generate.go

@@ -5,6 +5,7 @@ import (
 	"unicode"
 	"unicode/utf8"
 
+	"github.com/hashicorp/hcl2/hcl"
 	"github.com/hashicorp/hcl2/hcl/hclsyntax"
 	"github.com/zclconf/go-cty/cty"
 )
@@ -25,6 +26,19 @@ func TokensForValue(val cty.Value) Tokens {
 	return toks
 }
 
+// TokensForTraversal returns a sequence of tokens that represents the given
+// traversal.
+//
+// If the traversal is absolute then the result is a self-contained, valid
+// reference expression. If the traversal is relative then the returned tokens
+// could be appended to some other expression tokens to traverse into the
+// represented expression.
+func TokensForTraversal(traversal hcl.Traversal) Tokens {
+	toks := appendTokensForTraversal(traversal, nil)
+	format(toks) // fiddle with the SpacesBefore field to get canonical spacing
+	return toks
+}
+
 func appendTokensForValue(val cty.Value, toks Tokens) Tokens {
 	switch {
 
@@ -143,6 +157,47 @@ func appendTokensForValue(val cty.Value, toks Tokens) Tokens {
 	return toks
 }
 
+func appendTokensForTraversal(traversal hcl.Traversal, toks Tokens) Tokens {
+	for _, step := range traversal {
+		appendTokensForTraversalStep(step, toks)
+	}
+	return toks
+}
+
+func appendTokensForTraversalStep(step hcl.Traverser, toks Tokens) {
+	switch ts := step.(type) {
+	case hcl.TraverseRoot:
+		toks = append(toks, &Token{
+			Type:  hclsyntax.TokenIdent,
+			Bytes: []byte(ts.Name),
+		})
+	case hcl.TraverseAttr:
+		toks = append(
+			toks,
+			&Token{
+				Type:  hclsyntax.TokenDot,
+				Bytes: []byte{'.'},
+			},
+			&Token{
+				Type:  hclsyntax.TokenIdent,
+				Bytes: []byte(ts.Name),
+			},
+		)
+	case hcl.TraverseIndex:
+		toks = append(toks, &Token{
+			Type:  hclsyntax.TokenOBrack,
+			Bytes: []byte{'['},
+		})
+		appendTokensForValue(ts.Key, toks)
+		toks = append(toks, &Token{
+			Type:  hclsyntax.TokenCBrack,
+			Bytes: []byte{']'},
+		})
+	default:
+		panic(fmt.Sprintf("unsupported traversal step type %T", step))
+	}
+}
+
 func escapeQuotedStringLit(s string) []byte {
 	if len(s) == 0 {
 		return nil