From 2a6e423d2df023723acb487ca1bdec8f0d42630e Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Wed, 6 Jan 2021 20:44:28 +1100
Subject: [PATCH] Can assign-update tag

---
 pkg/yqlib/doc/Tag.md           | 16 +++++++++++++++-
 pkg/yqlib/operator_tag.go      | 28 +++++++++++++++++++---------
 pkg/yqlib/operator_tag_test.go | 18 +++++++++++++++++-
 3 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/pkg/yqlib/doc/Tag.md b/pkg/yqlib/doc/Tag.md
index 519142fd..ab64ed36 100644
--- a/pkg/yqlib/doc/Tag.md
+++ b/pkg/yqlib/doc/Tag.md
@@ -22,7 +22,21 @@ will output
 !!seq
 ```
 
-## Convert numbers to strings
+## Set custom tag
+Given a sample.yml file of:
+```yaml
+a: str
+```
+then
+```bash
+yq eval '.a tag = "!!mikefarah"' sample.yml
+```
+will output
+```yaml
+a: !!mikefarah str
+```
+
+## Find numbers and convert them to strings
 Given a sample.yml file of:
 ```yaml
 a: cat
diff --git a/pkg/yqlib/operator_tag.go b/pkg/yqlib/operator_tag.go
index 3ca2ede6..c9a9099e 100644
--- a/pkg/yqlib/operator_tag.go
+++ b/pkg/yqlib/operator_tag.go
@@ -9,15 +9,17 @@ import (
 func AssignTagOperator(d *dataTreeNavigator, matchingNodes *list.List, pathNode *PathTreeNode) (*list.List, error) {
 
 	log.Debugf("AssignTagOperator: %v")
-
-	rhs, err := d.GetMatchingNodes(matchingNodes, pathNode.Rhs)
-	if err != nil {
-		return nil, err
-	}
 	tag := ""
 
-	if rhs.Front() != nil {
-		tag = rhs.Front().Value.(*CandidateNode).Node.Value
+	if !pathNode.Operation.UpdateAssign {
+		rhs, err := d.GetMatchingNodes(matchingNodes, pathNode.Rhs)
+		if err != nil {
+			return nil, err
+		}
+
+		if rhs.Front() != nil {
+			tag = rhs.Front().Value.(*CandidateNode).Node.Value
+		}
 	}
 
 	lhs, err := d.GetMatchingNodes(matchingNodes, pathNode.Lhs)
@@ -26,11 +28,19 @@ func AssignTagOperator(d *dataTreeNavigator, matchingNodes *list.List, pathNode
 		return nil, err
 	}
 
-	// FAIL HERE
-
 	for el := lhs.Front(); el != nil; el = el.Next() {
 		candidate := el.Value.(*CandidateNode)
 		log.Debugf("Setting tag of : %v", candidate.GetKey())
+		if pathNode.Operation.UpdateAssign {
+			rhs, err := d.GetMatchingNodes(nodeToMap(candidate), pathNode.Rhs)
+			if err != nil {
+				return nil, err
+			}
+
+			if rhs.Front() != nil {
+				tag = rhs.Front().Value.(*CandidateNode).Node.Value
+			}
+		}
 		candidate.Node.Tag = tag
 	}
 
diff --git a/pkg/yqlib/operator_tag_test.go b/pkg/yqlib/operator_tag_test.go
index df47335b..b64afadd 100644
--- a/pkg/yqlib/operator_tag_test.go
+++ b/pkg/yqlib/operator_tag_test.go
@@ -27,13 +27,29 @@ var tagOperatorScenarios = []expressionScenario{
 		},
 	},
 	{
-		description: "Convert numbers to strings",
+		description: "Set custom tag",
+		document:    `{a: str}`,
+		expression:  `.a tag = "!!mikefarah"`,
+		expected: []string{
+			"D0, P[], (doc)::{a: !!mikefarah str}\n",
+		},
+	},
+	{
+		description: "Find numbers and convert them to strings",
 		document:    `{a: cat, b: 5, c: 3.2, e: true}`,
 		expression:  `(.. | select(tag == "!!int")) tag= "!!str"`,
 		expected: []string{
 			"D0, P[], (!!map)::{a: cat, b: \"5\", c: 3.2, e: true}\n",
 		},
 	},
+	{
+		skipDoc:    true,
+		document:   `{a: "!!frog", b: "!!customTag"}`,
+		expression: `.[] tag |= .`,
+		expected: []string{
+			"D0, P[], (doc)::{a: !!frog \"!!frog\", b: !!customTag \"!!customTag\"}\n",
+		},
+	},
 }
 
 func TestTagOperatorScenarios(t *testing.T) {