From 23083ed974223f3565194452789a057ac2dd1245 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Fri, 9 Oct 2020 17:07:53 +1100
Subject: [PATCH] fixed equals number issue

---
 pkg/yqlib/treeops/data_tree_navigator_test.go | 28 ++++++++++++++++
 pkg/yqlib/treeops/path_postfix_test.go        | 32 +++++++++++++++++++
 pkg/yqlib/treeops/path_tokeniser.go           |  4 +--
 pkg/yqlib/treeops/path_tokeniser_test.go      |  2 ++
 4 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/pkg/yqlib/treeops/data_tree_navigator_test.go b/pkg/yqlib/treeops/data_tree_navigator_test.go
index 55585ed..0ffc9f3 100644
--- a/pkg/yqlib/treeops/data_tree_navigator_test.go
+++ b/pkg/yqlib/treeops/data_tree_navigator_test.go
@@ -390,6 +390,34 @@ func TestDataTreeNavigatorArrayEquals(t *testing.T) {
 	test.AssertResult(t, expected, resultsToString(results))
 }
 
+func TestDataTreeNavigatorArrayEqualsDeep(t *testing.T) {
+
+	nodes := readDoc(t, `apples:
+  - { b: apple, animal: {legs: 2} }
+  - { b: banana, animal: {legs: 4} }
+  - { b: corn, animal: {legs: 6} }
+`)
+
+	path, errPath := treeCreator.ParsePath("apples(animal.legs == 4)")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [apples 1]
+  Tag: !!map, Kind: MappingNode, Anchor: 
+  {b: banana, animal: {legs: 4}}
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
 func TestDataTreeNavigatorEqualsTrickey(t *testing.T) {
 
 	nodes := readDoc(t, `a: 
diff --git a/pkg/yqlib/treeops/path_postfix_test.go b/pkg/yqlib/treeops/path_postfix_test.go
index 2e9d636..220ef6a 100644
--- a/pkg/yqlib/treeops/path_postfix_test.go
+++ b/pkg/yqlib/treeops/path_postfix_test.go
@@ -25,6 +25,20 @@ func testExpression(expression string) (string, error) {
 	return formatted, nil
 }
 
+func TestPostFixArrayEquals(t *testing.T) {
+	var infix = "a"
+	var expectedOutput = `PathKey - 'a'
+--------
+`
+
+	actual, err := testExpression(infix)
+	if err != nil {
+		t.Error(err)
+	}
+
+	test.AssertResultComplex(t, expectedOutput, actual)
+}
+
 func TestPostFixSimpleExample(t *testing.T) {
 	var infix = "a"
 	var expectedOutput = `PathKey - 'a'
@@ -167,6 +181,24 @@ Operation - OR
 	test.AssertResultComplex(t, expectedOutput, actual)
 }
 
+func TestPostFixEqualsNumberExample(t *testing.T) {
+	var infix = "(animal == 3)"
+	var expectedOutput = `PathKey - 'animal'
+--------
+PathKey - '3'
+--------
+Operation - EQUALS
+--------
+`
+
+	actual, err := testExpression(infix)
+	if err != nil {
+		t.Error(err)
+	}
+
+	test.AssertResultComplex(t, expectedOutput, actual)
+}
+
 func TestPostFixOrWithEqualsExample(t *testing.T) {
 	var infix = "a==thing OR b==thongs"
 	var expectedOutput = `PathKey - 'a'
diff --git a/pkg/yqlib/treeops/path_tokeniser.go b/pkg/yqlib/treeops/path_tokeniser.go
index acabe79..e1c42ec 100644
--- a/pkg/yqlib/treeops/path_tokeniser.go
+++ b/pkg/yqlib/treeops/path_tokeniser.go
@@ -142,14 +142,14 @@ func (p *pathTokeniser) Tokenise(path string) ([]*lex.Token, error) {
 
 	for index, token := range tokens {
 		for _, literalTokenDef := range append(Literals, "ARRAY_INDEX", "(") {
-			if index > 0 && token.Type == TokenIds[literalTokenDef] && tokens[index-1].Type != TokenIds["TRAVERSE_OPERATOR"] {
+			if index > 0 && token.Type == TokenIds[literalTokenDef] && tokens[index-1].Type != TokenIds["TRAVERSE_OPERATOR"] && tokens[index-1].Type != TokenIds["EQUALS_OPERATOR"] && tokens[index-1].Type != TokenIds["EQUALS_SELF_OPERATOR"] {
 				postProcessedTokens = append(postProcessedTokens, &lex.Token{Type: TokenIds["TRAVERSE_OPERATOR"], Value: "."})
 			}
 		}
 
 		postProcessedTokens = append(postProcessedTokens, token)
 		for _, literalTokenDef := range append(Literals, "ARRAY_INDEX", ")") {
-			if index != len(tokens)-1 && token.Type == TokenIds[literalTokenDef] && tokens[index+1].Type != TokenIds["TRAVERSE_OPERATOR"] {
+			if index != len(tokens)-1 && token.Type == TokenIds[literalTokenDef] && tokens[index+1].Type != TokenIds["TRAVERSE_OPERATOR"] && tokens[index+1].Type != TokenIds[")"] {
 				postProcessedTokens = append(postProcessedTokens, &lex.Token{Type: TokenIds["TRAVERSE_OPERATOR"], Value: "."})
 			}
 		}
diff --git a/pkg/yqlib/treeops/path_tokeniser_test.go b/pkg/yqlib/treeops/path_tokeniser_test.go
index 4e7405e..7c673f3 100644
--- a/pkg/yqlib/treeops/path_tokeniser_test.go
+++ b/pkg/yqlib/treeops/path_tokeniser_test.go
@@ -11,6 +11,8 @@ var tokeniserTests = []struct {
 	expectedTokens []interface{}
 }{ // TODO: Ensure ALL documented examples have tests! sheesh
 
+	{"(animal==3)", append(make([]interface{}, 0), "(", "animal", "==", int64(3), ")")},
+	{"(animal==f3)", append(make([]interface{}, 0), "(", "animal", "==", "f3", ")")},
 	{"apples.BANANAS", append(make([]interface{}, 0), "apples", ".", "BANANAS")},
 	{"appl*.BANA*", append(make([]interface{}, 0), "appl*", ".", "BANA*")},
 	{"a.b.**", append(make([]interface{}, 0), "a", ".", "b", ".", "**")},