From d04d1f2bd27a9c15293911e6a7023162c60fe227 Mon Sep 17 00:00:00 2001
From: Masayuki Morita <minamijoyo@gmail.com>
Date: Fri, 30 Aug 2024 17:36:31 +0900
Subject: [PATCH] Fix attribute get --with-comments for inline comments

Follow-up to #103.
Inline comments belong to attributes, not expressions. When the
`--with-comments` flag is specified, we need to return tokens on the
right side of the equals sign.
---
 cmd/attribute_test.go             | 14 +++++-----
 editor/sink_attribute_get.go      | 21 ++++++++++++---
 editor/sink_attribute_get_test.go | 44 +++++++++++++++++++++++++++++--
 3 files changed, 66 insertions(+), 13 deletions(-)

diff --git a/cmd/attribute_test.go b/cmd/attribute_test.go
index 41acb20..ada65c7 100644
--- a/cmd/attribute_test.go
+++ b/cmd/attribute_test.go
@@ -71,13 +71,13 @@ locals {
   }
 `,
 		},
-		// does not pass at current
-		// {
-		// 	name: "single with comments",
-		// 	args: []string{"--with-comments", "locals.attribute"},
-		// 	ok:   true,
-		// 	want: `"foo" #comment`,
-		// },
+		{
+			name: "single with comments",
+			args: []string{"--with-comments", "locals.attribute"},
+			ok:   true,
+			want: `"foo" # comment
+`,
+		},
 		{
 			name: "single without comments",
 			args: []string{"locals.attribute"},
diff --git a/editor/sink_attribute_get.go b/editor/sink_attribute_get.go
index 0be2827..6396681 100644
--- a/editor/sink_attribute_get.go
+++ b/editor/sink_attribute_get.go
@@ -186,13 +186,26 @@ func longestMatchingLabels(labels []string, prefix []string) []string {
 // There is no way to get value as string directly,
 // so we parses tokens of Attribute and build string representation.
 func GetAttributeValueAsString(attr *hclwrite.Attribute, withComments bool) (string, error) {
-	// find TokenEqual
-	expr := attr.Expr()
-	exprTokens := expr.BuildTokens(nil)
+	var rhsTokens hclwrite.Tokens
+	if withComments {
+		// Inline comments belong to an attribute, not an expression.
+		attrTokens := attr.BuildTokens(nil)
+		for i, t := range attrTokens {
+			// find TokenEqual
+			if t.Type != hclsyntax.TokenEqual {
+				continue
+			}
+			rhsTokens = attrTokens[i+1:]
+			break
+		}
+	} else {
+		expr := attr.Expr()
+		rhsTokens = expr.BuildTokens(nil)
+	}
 
 	// append tokens until find TokenComment
 	var valueTokens hclwrite.Tokens
-	for _, t := range exprTokens {
+	for _, t := range rhsTokens {
 		if t.Type == hclsyntax.TokenComment && !withComments {
 			t.Bytes = []byte("\n")
 			t.SpacesBefore = 0
diff --git a/editor/sink_attribute_get_test.go b/editor/sink_attribute_get_test.go
index fce8c55..672a7ab 100644
--- a/editor/sink_attribute_get_test.go
+++ b/editor/sink_attribute_get_test.go
@@ -48,7 +48,7 @@ a1 = v1
 			want:         "",
 		},
 		{
-			name: "attribute with comments",
+			name: "attribute without comments",
 			src: `
 // attr comment
 a0 = v0 // inline comment
@@ -60,7 +60,19 @@ a1 = v1
 			want:         "v0\n",
 		},
 		{
-			name: "multiline attribute with comments",
+			name: "attribute with comments",
+			src: `
+// attr comment
+a0 = v0 // inline comment
+a1 = v1
+`,
+			address:      "a0",
+			withComments: true,
+			ok:           true,
+			want:         "v0 // inline comment\n",
+		},
+		{
+			name: "multiline attribute without comments",
 			src: `
 // attr comment
 a0 = v0
@@ -85,6 +97,34 @@ a2 = v2
 
   "val5",
 ]
+`,
+		},
+		{
+			name: "multiline attribute with comments",
+			src: `
+// attr comment
+a0 = v0
+a1 = [
+  "val1",
+  "val2", // inline comment
+  "val3", # another comment
+  "val4",
+  # a ocmment line
+  "val5",
+]
+a2 = v2
+`,
+			address:      "a1",
+			withComments: true,
+			ok:           true,
+			want: `[
+  "val1",
+  "val2", // inline comment
+  "val3", # another comment
+  "val4",
+  # a ocmment line
+  "val5",
+]
 `,
 		},
 		{