hclwrite: Rewrite comments on blocks and attributes

hclwrite/ast_block.go

@@ -96,6 +96,43 @@ func (b *Block) SetLabels(labels []string) {
 	b.labelsObj().Replace(labels)
 }
 
+// Comments returns the comments that annotate the block.
+//
+// The return value is an empty slice if the block has no lead comments, or a
+// slice of strings representing the raw text of the comments, including
+// leading and trailing comment markers.
+func (b *Block) Comments() []string {
+	tokens := b.leadComments.content.(*comments).tokens
+	blockComments := make([]string, len(tokens))
+	for i := range tokens {
+		blockComments[i] = string(tokens[i].Bytes)
+	}
+	return blockComments
+}
+
+// SetComments replaces the comments that annotate the block.
+//
+// Each item should contain any leading and trailing markers, i.e. single-
+// line comments should start with either “#” or “//” and end with a
+// newline character, and block comments should start with “/*” and end with
+// “*/”.
+func (b *Block) SetComments(blockComments []string) {
+	c := b.leadComments.content.(*comments)
+	c.tokens = make(Tokens, len(blockComments))
+
+	for i := range blockComments {
+		c.tokens[i] = &Token{
+			Type:  hclsyntax.TokenComment,
+			Bytes: []byte(blockComments[i]),
+		}
+	}
+}
+
+// Single-comment alternative to [SetComments]
+func (b *Block) SetComment(blockComment string) {
+	b.SetComments([]string{blockComment})
+}
+
 // labelsObj returns the internal node content representation of the block
 // labels. This is not part of the public API because we're intentionally
 // exposing only a limited API to get/set labels on the block itself in a

hclwrite/ast_block_test.go

@@ -487,3 +487,182 @@ func TestBlockSetLabels(t *testing.T) {
 		})
 	}
 }
+
+func TestBlockComments(t *testing.T) {
+	tests := []struct {
+		src  string
+		want []string
+	}{
+		{
+			`
+block {
+}
+`,
+			[]string{},
+		},
+		{
+			`
+# Comment
+block {
+}
+`,
+			[]string{"# Comment\n"},
+		},
+		{
+			`
+# First line
+# Second line
+block {
+}
+`,
+			[]string{"# First line\n", "# Second line\n"},
+		},
+		{
+			`
+/* Comment */ block {
+}
+`,
+			[]string{"/* Comment */"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%s", strings.Join(test.want, " ")), 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")
+			}
+
+			block := f.Body().Blocks()[0]
+			got := block.Comments()
+			if !reflect.DeepEqual(got, test.want) {
+				t.Errorf("wrong result\ngot:  %#v\nwant: %#v", got, test.want)
+			}
+		})
+	}
+}
+
+func TestBlockSetComments(t *testing.T) {
+	tests := []struct {
+		src           string
+		typeName      string
+		blockComments []string
+		want          Tokens
+	}{
+		{
+			`block {}`,
+			"block",
+			[]string{"# Comment\n"},
+			Tokens{
+				{
+					Type:         hclsyntax.TokenComment,
+					Bytes:        []byte("# Comment\n"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte(`block`),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenOBrace,
+					Bytes:        []byte{'{'},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenCBrace,
+					Bytes:        []byte{'}'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+		{
+			"# Comment\nblock {}",
+			"block",
+			nil,
+			Tokens{
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte(`block`),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenOBrace,
+					Bytes:        []byte{'{'},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenCBrace,
+					Bytes:        []byte{'}'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+		{
+			"# Old\nblock {}",
+			"block",
+			[]string{"# New\n"},
+			Tokens{
+				{
+					Type:         hclsyntax.TokenComment,
+					Bytes:        []byte("# New\n"),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte(`block`),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenOBrace,
+					Bytes:        []byte{'{'},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenCBrace,
+					Bytes:        []byte{'}'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%s %s in %s", test.typeName, test.blockComments, 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")
+			}
+
+			b := f.Body().FirstMatchingBlock(test.typeName, nil)
+			b.SetComments(test.blockComments)
+			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)
+			}
+		})
+	}
+}

hclwrite/ast_body.go

@@ -137,6 +137,111 @@ func (b *Body) RemoveBlock(block *Block) bool {
 	return false
 }
 
+// AttributeLeadComments returns the comments that annotate an attribute.
+//
+// The return value is nil if there was no such attribute, an empty slice if
+// it had no lead comments, or a slice of strings representing the raw text of
+// the comments, including leading and trailing comment markers.
+func (b *Body) AttributeLeadComments(name string) []string {
+	attr := b.GetAttribute(name)
+	if attr == nil {
+		return nil
+	}
+
+	tokens := attr.leadComments.content.(*comments).tokens
+	attrComments := make([]string, len(tokens))
+	for i := range tokens {
+		attrComments[i] = string(tokens[i].Bytes)
+	}
+	return attrComments
+}
+
+// SetAttributeLeadComments replaces the comments that annotate an attribute.
+//
+// Each item should contain any leading and trailing markers, i.e. single-
+// line comments should start with either “#” or “//” and end with a
+// newline character, and block comments should start with “/*” and end with
+// “*/”.
+func (b *Body) SetAttributeLeadComments(name string, leadComments []string) *Attribute {
+	attr := b.GetAttribute(name)
+	if attr == nil {
+		return nil
+	}
+
+	c := attr.leadComments.content.(*comments)
+	c.tokens = make(Tokens, len(leadComments))
+
+	for i := range leadComments {
+		c.tokens[i] = &Token{
+			Type:  hclsyntax.TokenComment,
+			Bytes: []byte(leadComments[i]),
+		}
+	}
+
+	return attr
+}
+
+// Single-comment alternative to [SetAttributeLeadComments]
+func (b *Body) SetAttributeLeadComment(name string, leadComment string) *Attribute {
+	return b.SetAttributeLeadComments(name, []string{leadComment})
+}
+
+// AttributeLineComments returns the comments that follow an attribute on
+// the same line.
+//
+// The return value is nil if there was no such attribute, an empty slice if
+// it had no line comments, or a slice of strings representing the raw text of
+// the comments, including leading and trailing comment markers.
+func (b *Body) AttributeLineComments(name string) []string {
+	attr := b.GetAttribute(name)
+	if attr == nil {
+		return nil
+	}
+
+	tokens := attr.lineComments.content.(*comments).tokens
+	attrComments := make([]string, len(tokens))
+	for i := range tokens {
+		attrComments[i] = string(tokens[i].Bytes)
+	}
+	return attrComments
+}
+
+// SetAttributeLineComments replaces the comments that follow an attribute on
+// the same line.
+//
+// Each item should contain any leading and trailing markers, i.e. single-
+// line comments should start with either “#” or “//” and end with a
+// newline character, and block comments should start with “/*” and end with
+// “*/”.
+//
+// While it is technically possible to set multiple single-line comments using
+// this function, doing so is discouraged since the parser would never produce
+// such constructs. For single-line comments, prefer to use
+// [SetAttributeLineComment].
+func (b *Body) SetAttributeLineComments(name string, lineComments []string) *Attribute {
+	attr := b.GetAttribute(name)
+	if attr == nil {
+		return nil
+	}
+
+	c := attr.lineComments.content.(*comments)
+	c.tokens = make(Tokens, len(lineComments))
+
+	for i := range lineComments {
+		c.tokens[i] = &Token{
+			Type:  hclsyntax.TokenComment,
+			Bytes: []byte(lineComments[i]),
+		}
+	}
+
+	return attr
+}
+
+// Single-comment alternative to [SetAttributeLineComments]
+func (b *Body) SetAttributeLineComment(name string, lineComment string) *Attribute {
+	return b.SetAttributeLineComments(name, []string{lineComment})
+}
+
 // SetAttributeRaw either replaces the expression of an existing attribute
 // of the given name or adds a new attribute definition to the end of the block,
 // using the given tokens verbatim as the expression.