Merge pull request #369 from bendrucker/hclwriter-numeric-dot

hclwrite: handle legacy dot access of numeric indexes

2 files changed, 171 insertions, 0 deletions

diff --git a/hclwrite/parser.go b/hclwrite/parser.go
index d6cf532..8e100fb 100644
--- a/hclwrite/parser.go
+++ b/hclwrite/parser.go
@@ -88,6 +88,16 @@ func (it inputTokens) PartitionType(ty hclsyntax.TokenType) (before, within, aft
 	panic(fmt.Sprintf("didn't find any token of type %s", ty))
 }
 
+func (it inputTokens) PartitionTypeOk(ty hclsyntax.TokenType) (before, within, after inputTokens, ok bool) {
+	for i, t := range it.writerTokens {
+		if t.Type == ty {
+			return it.Slice(0, i), it.Slice(i, i+1), it.Slice(i+1, len(it.nativeTokens)), true
+		}
+	}
+
+	return inputTokens{}, inputTokens{}, inputTokens{}, false
+}
+
 func (it inputTokens) PartitionTypeSingle(ty hclsyntax.TokenType) (before inputTokens, found *Token, after inputTokens) {
 	before, within, after := it.PartitionType(ty)
 	if within.Len() != 1 {
@@ -404,6 +414,19 @@ func parseTraversalStep(nativeStep hcl.Traverser, from inputTokens) (before inpu
 		children = step.inTree.children
 		before, from, after = from.Partition(nativeStep.SourceRange())
 
+		if inBefore, dot, from, ok := from.PartitionTypeOk(hclsyntax.TokenDot); ok {
+			children.AppendUnstructuredTokens(inBefore.Tokens())
+			children.AppendUnstructuredTokens(dot.Tokens())
+
+			valBefore, valToken, valAfter := from.PartitionTypeSingle(hclsyntax.TokenNumberLit)
+			children.AppendUnstructuredTokens(valBefore.Tokens())
+			key := newNumber(valToken)
+			step.key = children.Append(key)
+			children.AppendUnstructuredTokens(valAfter.Tokens())
+
+			return before, newNode(step), after
+		}
+
 		var inBefore, oBrack, keyTokens, cBrack inputTokens
 		inBefore, oBrack, from = from.PartitionType(hclsyntax.TokenOBrack)
 		children.AppendUnstructuredTokens(inBefore.Tokens())
diff --git a/hclwrite/parser_test.go b/hclwrite/parser_test.go
index 28b5380..9c57c6c 100644
--- a/hclwrite/parser_test.go
+++ b/hclwrite/parser_test.go
@@ -556,6 +556,69 @@ func TestParse(t *testing.T) {
 			},
 		},
 		{
+			"a = foo.0\n",
+			TestTreeNode{
+				Type: "Body",
+				Children: []TestTreeNode{
+					{
+						Type: "Attribute",
+						Children: []TestTreeNode{
+							{
+								Type: "comments",
+							},
+							{
+								Type: "identifier",
+								Val:  "a",
+							},
+							{
+								Type: "Tokens",
+								Val:  " =",
+							},
+							{
+								Type: "Expression",
+								Children: []TestTreeNode{
+									{
+										Type: "Traversal",
+										Children: []TestTreeNode{
+											{
+												Type: "TraverseName",
+												Children: []TestTreeNode{
+													{
+														Type: "identifier",
+														Val:  " foo",
+													},
+												},
+											},
+											{
+												Type: "TraverseIndex",
+												Children: []TestTreeNode{
+													{
+														Type: "Tokens",
+														Val:  ".",
+													},
+													{
+														Type: "number",
+														Val:  "0",
+													},
+												},
+											},
+										},
+									},
+								},
+							},
+							{
+								Type: "comments",
+							},
+							{
+								Type: "Tokens",
+								Val:  "\n",
+							},
+						},
+					},
+				},
+			},
+		},
+		{
 			"a = foo[bar]\n",
 			TestTreeNode{
 				Type: "Body",
@@ -628,6 +691,91 @@ func TestParse(t *testing.T) {
 			},
 		},
 		{
+			"a = foo[bar.baz]\n",
+			TestTreeNode{
+				Type: "Body",
+				Children: []TestTreeNode{
+					{
+						Type: "Attribute",
+						Children: []TestTreeNode{
+							{
+								Type: "comments",
+							},
+							{
+								Type: "identifier",
+								Val:  "a",
+							},
+							{
+								Type: "Tokens",
+								Val:  " =",
+							},
+							{
+								Type: "Expression",
+								Children: []TestTreeNode{
+									{
+										Type: "Traversal",
+										Children: []TestTreeNode{
+											{
+												Type: "TraverseName",
+												Children: []TestTreeNode{
+													{
+														Type: "identifier",
+														Val:  " foo",
+													},
+												},
+											},
+										},
+									},
+									{
+										Type: "Tokens",
+										Val:  "[",
+									},
+									{
+										Type: "Traversal",
+										Children: []TestTreeNode{
+											{
+												Type: "TraverseName",
+												Children: []TestTreeNode{
+													{
+														Type: "identifier",
+														Val:  "bar",
+													},
+												},
+											},
+											{
+												Type: "TraverseName",
+												Children: []TestTreeNode{
+													{
+														Type: "Tokens",
+														Val:  ".",
+													},
+													{
+														Type: "identifier",
+														Val:  "baz",
+													},
+												},
+											},
+										},
+									},
+									{
+										Type: "Tokens",
+										Val:  "]",
+									},
+								},
+							},
+							{
+								Type: "comments",
+							},
+							{
+								Type: "Tokens",
+								Val:  "\n",
+							},
+						},
+					},
+				},
+			},
+		},
+		{
 			"a = foo[bar].baz\n",
 			TestTreeNode{
 				Type: "Body",