support string X int multiplication as jq (#1988)

mikefarah · Mar 22, 2024 · 3b85cef · 3b85cef
1 parent 76a0da3
commit 3b85cef
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 0 deletions.
diff --git a/pkg/yqlib/doc/operators/multiply-merge.md b/pkg/yqlib/doc/operators/multiply-merge.md
@@ -55,6 +55,62 @@ a: 12
 b: 4
 ```
 
+## Multiply string node X int
+Given a sample.yml file of:
+```yaml
+b: banana
+```
+then
+```bash
+yq '.b * 4' sample.yml
+```
+will output
+```yaml
+bananabananabananabanana
+```
+
+## Multiply int X string node
+Given a sample.yml file of:
+```yaml
+b: banana
+```
+then
+```bash
+yq '4 * .b' sample.yml
+```
+will output
+```yaml
+bananabananabananabanana
+```
+
+## Multiply string X int node
+Given a sample.yml file of:
+```yaml
+n: 4
+```
+then
+```bash
+yq '"banana" * .n' sample.yml
+```
+will output
+```yaml
+bananabananabananabanana
+```
+
+## Multiply int node X string
+Given a sample.yml file of:
+```yaml
+n: 4
+```
+then
+```bash
+yq '.n * "banana"' sample.yml
+```
+will output
+```yaml
+bananabananabananabanana
+```
+
 ## Merge objects together, returning merged result only
 Given a sample.yml file of:
 ```yaml

diff --git a/pkg/yqlib/operator_multiply.go b/pkg/yqlib/operator_multiply.go
@@ -91,6 +91,8 @@ func multiplyScalars(lhs *CandidateNode, rhs *CandidateNode) (*CandidateNode, er
 		return multiplyIntegers(lhs, rhs)
 	} else if (lhsTag == "!!int" || lhsTag == "!!float") && (rhsTag == "!!int" || rhsTag == "!!float") {
 		return multiplyFloats(lhs, rhs, lhsIsCustom)
+	} else if (lhsTag == "!!str" && rhsTag == "!!int") || (lhsTag == "!!int" && rhsTag == "!!str") {
+		return repeatString(lhs, rhs)
 	}
 	return nil, fmt.Errorf("cannot multiply %v with %v", lhs.Tag, rhs.Tag)
 }
@@ -135,6 +137,28 @@ func multiplyIntegers(lhs *CandidateNode, rhs *CandidateNode) (*CandidateNode, e
 	return target, nil
 }
 
+func repeatString(lhs *CandidateNode, rhs *CandidateNode) (*CandidateNode, error) {
+	var stringNode *CandidateNode
+	var intNode *CandidateNode
+	if lhs.Tag == "!!str" {
+		stringNode = lhs
+		intNode = rhs
+	} else {
+		stringNode = rhs
+		intNode = lhs
+	}
+	target := lhs.CopyWithoutContent()
+	target.UpdateAttributesFrom(stringNode, assignPreferences{})
+
+	count, err := parseInt(intNode.Value)
+	if err != nil {
+		return nil, err
+	}
+	target.Value = strings.Repeat(stringNode.Value, count)
+
+	return target, nil
+}
+
 func mergeObjects(d *dataTreeNavigator, context Context, lhs *CandidateNode, rhs *CandidateNode, preferences multiplyPreferences) (*CandidateNode, error) {
 	var results = list.New()
 

diff --git a/pkg/yqlib/operator_multiply_test.go b/pkg/yqlib/operator_multiply_test.go
@@ -1,6 +1,8 @@
 package yqlib
 
 import (
+	"fmt"
+	"strings"
 	"testing"
 )
 
@@ -174,6 +176,40 @@ var multiplyOperatorScenarios = []expressionScenario{
 			"D0, P[], (!!map)::a: 12\nb: 4\n",
 		},
 	},
+	{
+		description: "Multiply string node X int",
+		document:    docNoComments,
+		expression:  ".b * 4",
+		expected: []string{
+			fmt.Sprintf("D0, P[b], (!!str)::%s\n", strings.Repeat("banana", 4)),
+		},
+	},
+	{
+		description: "Multiply int X string node",
+		document:    docNoComments,
+		expression:  "4 * .b",
+		expected: []string{
+			fmt.Sprintf("D0, P[], (!!str)::%s\n", strings.Repeat("banana", 4)),
+		},
+	},
+	{
+		description: "Multiply string X int node",
+		document: `n: 4
+`,
+		expression: `"banana" * .n`,
+		expected: []string{
+			fmt.Sprintf("D0, P[], (!!str)::%s\n", strings.Repeat("banana", 4)),
+		},
+	},
+	{
+		description: "Multiply int node X string",
+		document: `n: 4
+`,
+		expression: `.n * "banana"`,
+		expected: []string{
+			fmt.Sprintf("D0, P[n], (!!str)::%s\n", strings.Repeat("banana", 4)),
+		},
+	},
 	{
 		skipDoc:    true,
 		document:   doc1,