add --with-comments to preserve comments when returning an attribute

minamijoyo · Aug 27, 2024 · 87c1cdb · 87c1cdb
1 parent 26c04bc
commit 87c1cdb
Show file tree

Hide file tree

Showing 6 changed files with 172 additions and 98 deletions.
diff --git a/cmd/attribute.go b/cmd/attribute.go
@@ -44,6 +44,10 @@ Arguments:
 		RunE: runAttributeGetCmd,
 	}
 
+	flags := cmd.Flags()
+	flags.Bool("with-comments", false, "return comments along with attribute value")
+	_ = viper.BindPFlag("attribute.get.withcomments", flags.Lookup("with-comments"))
+
 	return cmd
 }
 
@@ -55,11 +59,12 @@ func runAttributeGetCmd(cmd *cobra.Command, args []string) error {
 	address := args[0]
 	file := viper.GetString("file")
 	update := viper.GetBool("update")
+	withcomments := viper.GetBool("attribute.get.withcomments")
 	if update {
 		return errors.New("The update flag is not allowed")
 	}
 
-	sink := editor.NewAttributeGetSink(address)
+	sink := editor.NewAttributeGetSink(address, withcomments)
 	c := newDefaultClient(cmd)
 	return c.Derive(file, sink)
 }

diff --git a/cmd/attribute_test.go b/cmd/attribute_test.go
@@ -12,6 +12,13 @@ func TestAttributeGet(t *testing.T) {
     key    = "services/hoge/dev/terraform.tfstate"
   }
 }
+locals {
+  map = {
+    # comment
+    attribute = "bar"
+  }
+  attribute = "foo" # comment
+}
 `
 
 	cases := []struct {
@@ -44,6 +51,40 @@ func TestAttributeGet(t *testing.T) {
 			ok:   false,
 			want: "",
 		},
+		{
+			name: "with comments",
+			args: []string{"--with-comments", "locals.map"},
+			ok:   true,
+			want: `{
+    # comment
+    attribute = "bar"
+  }
+`,
+		},
+		{
+			name: "without comments",
+			args: []string{"locals.map"},
+			ok:   true,
+			want: `{
+
+    attribute = "bar"
+  }
+`,
+		},
+		// does not pass at current
+		// {
+		// 	name: "single with comments",
+		// 	args: []string{"--with-comments", "locals.attribute"},
+		// 	ok:   true,
+		// 	want: `"foo" #comment`,
+		// },
+		{
+			name: "single without comments",
+			args: []string{"locals.attribute"},
+			ok:   true,
+			want: `"foo"
+`,
+		},
 	}
 
 	for _, tc := range cases {

diff --git a/editor/client_test.go b/editor/client_test.go
@@ -124,7 +124,7 @@ a1 = v1
 a0 = v3
 a2 = v2
 `
-	sink := NewAttributeGetSink("a0")
+	sink := NewAttributeGetSink("a0", false)
 
 	cases := []struct {
 		name     string

diff --git a/editor/operator_derive_test.go b/editor/operator_derive_test.go
@@ -7,31 +7,34 @@ import (
 
 func TestOperatorDeriveApply(t *testing.T) {
 	cases := []struct {
-		name    string
-		src     string
-		address string
-		ok      bool
-		want    string
+		name         string
+		src          string
+		address      string
+		withcomments bool
+		ok           bool
+		want         string
 	}{
 		{
 			name: "match",
 			src: `
 a0 = v0
 a1 = v1
 `,
-			address: "a0",
-			ok:      true,
-			want:    "v0\n",
+			address:      "a0",
+			withcomments: false,
+			ok:           true,
+			want:         "v0\n",
 		},
 		{
 			name: "not found",
 			src: `
 a0 = v0
 a1 = v1
 `,
-			address: "a2",
-			ok:      true,
-			want:    "",
+			address:      "a2",
+			withcomments: false,
+			ok:           true,
+			want:         "",
 		},
 		{
 			name: "syntax error",
@@ -46,7 +49,7 @@ b1 {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			o := NewDeriveOperator(NewAttributeGetSink(tc.address))
+			o := NewDeriveOperator(NewAttributeGetSink(tc.address, tc.withcomments))
 			output, err := o.Apply([]byte(tc.src), "test")
 			if tc.ok && err != nil {
 				t.Fatalf("unexpected err = %s", err)
@@ -66,39 +69,42 @@ b1 {
 
 func TestDeriveStream(t *testing.T) {
 	cases := []struct {
-		name    string
-		src     string
-		address string
-		ok      bool
-		want    string
+		name         string
+		src          string
+		address      string
+		withcomments bool
+		ok           bool
+		want         string
 	}{
 		{
 			name: "match",
 			src: `
 a0 = v0
 a1 = v1
 `,
-			address: "a0",
-			ok:      true,
-			want:    "v0\n",
+			address:      "a0",
+			withcomments: false,
+			ok:           true,
+			want:         "v0\n",
 		},
 		{
 			name: "not found",
 			src: `
 a0 = v0
 a1 = v1
 `,
-			address: "a2",
-			ok:      true,
-			want:    "",
+			address:      "a2",
+			withcomments: false,
+			ok:           true,
+			want:         "",
 		},
 	}
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
 			inStream := bytes.NewBufferString(tc.src)
 			outStream := new(bytes.Buffer)
-			sink := NewAttributeGetSink(tc.address)
+			sink := NewAttributeGetSink(tc.address, tc.withcomments)
 			err := DeriveStream(inStream, outStream, "test", sink)
 			if tc.ok && err != nil {
 				t.Fatalf("unexpected err = %s", err)
@@ -118,40 +124,43 @@ a1 = v1
 
 func TestDeriveFile(t *testing.T) {
 	cases := []struct {
-		name    string
-		src     string
-		address string
-		value   string
-		ok      bool
-		want    string
+		name         string
+		src          string
+		address      string
+		withcomments bool
+		value        string
+		ok           bool
+		want         string
 	}{
 		{
 			name: "match",
 			src: `
 a0 = v0
 a1 = v1
 `,
-			address: "a0",
-			ok:      true,
-			want:    "v0\n",
+			address:      "a0",
+			withcomments: false,
+			ok:           true,
+			want:         "v0\n",
 		},
 		{
 			name: "not found",
 			src: `
 a0 = v0
 a1 = v1
 `,
-			address: "a2",
-			ok:      true,
-			want:    "",
+			address:      "a2",
+			withcomments: false,
+			ok:           true,
+			want:         "",
 		},
 	}
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
 			path := setupTestFile(t, tc.src)
 			outStream := new(bytes.Buffer)
-			sink := NewAttributeGetSink(tc.address)
+			sink := NewAttributeGetSink(tc.address, tc.withcomments)
 			err := DeriveFile(path, outStream, sink)
 			if tc.ok && err != nil {
 				t.Fatalf("unexpected err = %s", err)
@@ -175,7 +184,7 @@ a1 = v1
 }
 
 func TestDeriveFileNotFound(t *testing.T) {
-	sink := NewAttributeGetSink("foo")
+	sink := NewAttributeGetSink("foo", false)
 	outStream := new(bytes.Buffer)
 	err := DeriveFile("not_found", outStream, sink)
 	if err == nil {

diff --git a/editor/sink_attribute_get.go b/editor/sink_attribute_get.go
@@ -11,14 +11,16 @@ import (
 // AttributeGetSink is a sink implementation for getting a value of attribute.
 type AttributeGetSink struct {
 	address string
+	withcomments bool
 }
 
 var _ Sink = (*AttributeGetSink)(nil)
 
 // NewAttributeGetSink creates a new instance of AttributeGetSink.
-func NewAttributeGetSink(address string) Sink {
+func NewAttributeGetSink(address string, withcomments bool) Sink {
 	return &AttributeGetSink{
 		address: address,
+		withcomments: withcomments,
 	}
 }
 
@@ -35,7 +37,7 @@ func (s *AttributeGetSink) Sink(inFile *hclwrite.File) ([]byte, error) {
 	}
 
 	// treat expr as a string without interpreting its meaning.
-	out, err := GetAttributeValueAsString(attr)
+	out, err := GetAttributeValueAsString(attr, s.withcomments)
 	if err != nil {
 		return []byte{}, err
 	}
@@ -183,15 +185,15 @@ func longestMatchingLabels(labels []string, prefix []string) []string {
 // GetAttributeValueAsString returns a value of Attribute as 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) (string, error) {
+func GetAttributeValueAsString(attr *hclwrite.Attribute, withcomments bool) (string, error) {
 	// find TokenEqual
 	expr := attr.Expr()
 	exprTokens := expr.BuildTokens(nil)
-
+	
 	// append tokens until find TokenComment
 	var valueTokens hclwrite.Tokens
 	for _, t := range exprTokens {
-		if t.Type == hclsyntax.TokenComment {
+		if t.Type == hclsyntax.TokenComment && !withcomments {
 			t.Bytes = []byte("\n")
 			t.SpacesBefore = 0
 		}