From c1b58f71997308af769b18083b3e78d9bdb5cd19 Mon Sep 17 00:00:00 2001 From: TJ Miller Date: Thu, 9 Mar 2023 13:22:45 -0800 Subject: [PATCH] Added divide operator (#49) --- pkg/yqlib/doc/operators/divide.md | 38 ++++++++++ pkg/yqlib/lexer_participle.go | 2 + pkg/yqlib/lib.go | 2 + pkg/yqlib/operator_divide.go | 111 +++++++++++++++++++++++++++ pkg/yqlib/operator_divide_test.go | 121 ++++++++++++++++++++++++++++++ 5 files changed, 274 insertions(+) create mode 100644 pkg/yqlib/doc/operators/divide.md create mode 100644 pkg/yqlib/operator_divide.go create mode 100644 pkg/yqlib/operator_divide_test.go diff --git a/pkg/yqlib/doc/operators/divide.md b/pkg/yqlib/doc/operators/divide.md new file mode 100644 index 00000000..bf005e4a --- /dev/null +++ b/pkg/yqlib/doc/operators/divide.md @@ -0,0 +1,38 @@ + +## String split +Given a sample.yml file of: +```yaml +a: cat_meow +b: _ +``` +then +```bash +yq '.c = .a / .b' sample.yml +``` +will output +```yaml +a: cat_meow +b: _ +c: + - cat + - meow +``` + +## Number division +The result during divison is calculated as a float + +Given a sample.yml file of: +```yaml +a: 12 +b: 2.5 +``` +then +```bash +yq '.a = .a / .b' sample.yml +``` +will output +```yaml +a: 4.8 +b: 2.5 +``` + diff --git a/pkg/yqlib/lexer_participle.go b/pkg/yqlib/lexer_participle.go index e6d4d58f..a63f8a79 100644 --- a/pkg/yqlib/lexer_participle.go +++ b/pkg/yqlib/lexer_participle.go @@ -210,6 +210,8 @@ var participleYqRules = []*participleYqRule{ {"MultiplyAssign", `\*=[\+|\?cdn]*`, multiplyWithPrefs(multiplyAssignOpType), 0}, {"Multiply", `\*[\+|\?cdn]*`, multiplyWithPrefs(multiplyOpType), 0}, + {"Divide", `\/`, opToken(divideOpType), 0}, + {"AddAssign", `\+=`, opToken(addAssignOpType), 0}, {"Add", `\+`, opToken(addOpType), 0}, diff --git a/pkg/yqlib/lib.go b/pkg/yqlib/lib.go index a24d18f2..f8e8957a 100644 --- a/pkg/yqlib/lib.go +++ b/pkg/yqlib/lib.go @@ -62,6 +62,8 @@ var assignAliasOpType = &operationType{Type: "ASSIGN_ALIAS", NumArgs: 2, Precede var multiplyOpType = &operationType{Type: "MULTIPLY", NumArgs: 2, Precedence: 42, Handler: multiplyOperator} var multiplyAssignOpType = &operationType{Type: "MULTIPLY_ASSIGN", NumArgs: 2, Precedence: 42, Handler: multiplyAssignOperator} +var divideOpType = &operationType{Type: "DIVIDE", NumArgs: 2, Precedence: 42, Handler: divideOperator} + var addOpType = &operationType{Type: "ADD", NumArgs: 2, Precedence: 42, Handler: addOperator} var subtractOpType = &operationType{Type: "SUBTRACT", NumArgs: 2, Precedence: 42, Handler: subtractOperator} var alternativeOpType = &operationType{Type: "ALTERNATIVE", NumArgs: 2, Precedence: 42, Handler: alternativeOperator} diff --git a/pkg/yqlib/operator_divide.go b/pkg/yqlib/operator_divide.go new file mode 100644 index 00000000..67cfb7da --- /dev/null +++ b/pkg/yqlib/operator_divide.go @@ -0,0 +1,111 @@ +package yqlib + +import ( + "fmt" + "strconv" + "strings" + + yaml "gopkg.in/yaml.v3" +) + +func createDivideOp(lhs *ExpressionNode, rhs *ExpressionNode) *ExpressionNode { + return &ExpressionNode{Operation: &Operation{OperationType: divideOpType}, + LHS: lhs, + RHS: rhs} +} + +// func divideAssignOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) { +// return compoundAssignFunction(d, context, expressionNode, createDivideOp) +// } + +func divideOperator(d *dataTreeNavigator, context Context, expressionNode *ExpressionNode) (Context, error) { + log.Debugf("Divide operator") + + return crossFunction(d, context.ReadOnlyClone(), expressionNode, divide, false) +} + +func divide(d *dataTreeNavigator, context Context, lhs *CandidateNode, rhs *CandidateNode) (*CandidateNode, error) { + lhs.Node = unwrapDoc(lhs.Node) + rhs.Node = unwrapDoc(rhs.Node) + + lhsNode := lhs.Node + + if lhsNode.Tag == "!!null" { + return nil, fmt.Errorf("%v (%v) cannot be divided by %v (%v)", lhsNode.Tag, lhs.GetNicePath(), rhs.Node.Tag, rhs.GetNicePath()) + } + + target := lhs.CreateReplacement(&yaml.Node{ + Anchor: lhs.Node.Anchor, + }) + + switch lhsNode.Kind { + case yaml.MappingNode: + return nil, fmt.Errorf("%v (%v) cannot be divided by %v (%v)", lhsNode.Tag, lhs.GetNicePath(), rhs.Node.Tag, rhs.GetNicePath()) + + case yaml.SequenceNode: + return nil, fmt.Errorf("%v (%v) cannot be divided by %v (%v)", lhsNode.Tag, lhs.GetNicePath(), rhs.Node.Tag, rhs.GetNicePath()) + + case yaml.ScalarNode: + if rhs.Node.Kind != yaml.ScalarNode { + return nil, fmt.Errorf("%v (%v) cannot be divided by %v (%v)", lhsNode.Tag, lhs.GetNicePath(), rhs.Node.Tag, rhs.GetNicePath()) + } + if err := divideScalars(context, target, lhsNode, rhs.Node); err != nil { + return nil, err + } + } + return target, nil +} + +func divideScalars(context Context, target *CandidateNode, lhs *yaml.Node, rhs *yaml.Node) error { + lhsTag := lhs.Tag + rhsTag := guessTagFromCustomType(rhs) + lhsIsCustom := false + if !strings.HasPrefix(lhsTag, "!!") { + // custom tag - we have to have a guess + lhsTag = guessTagFromCustomType(lhs) + lhsIsCustom = true + } + + if lhsTag == "!!str" && rhsTag == "!!str" { + target.Node = split(lhs.Value, rhs.Value) + target.Node.Anchor = lhs.Anchor + } else if lhsTag == "!!int" && rhsTag == "!!int" { + target.Node.Kind = yaml.ScalarNode + target.Node.Style = lhs.Style + + format, lhsNum, err := parseInt64(lhs.Value) + if err != nil { + return err + } + _, rhsNum, err := parseInt64(rhs.Value) + if err != nil { + return err + } + quotient := float64(lhsNum) / float64(rhsNum) + + target.Node.Tag = "!!float" + target.Node.Value = fmt.Sprintf(format, quotient) + } else if (lhsTag == "!!int" || lhsTag == "!!float") && (rhsTag == "!!int" || rhsTag == "!!float") { + target.Node.Kind = yaml.ScalarNode + target.Node.Style = lhs.Style + + lhsNum, err := strconv.ParseFloat(lhs.Value, 64) + if err != nil { + return err + } + rhsNum, err := strconv.ParseFloat(rhs.Value, 64) + if err != nil { + return err + } + quotient := lhsNum / rhsNum + if lhsIsCustom { + target.Node.Tag = lhs.Tag + } else { + target.Node.Tag = "!!float" + } + target.Node.Value = fmt.Sprintf("%v", quotient) + } else { + return fmt.Errorf("%v cannot be divided by %v", lhsTag, rhsTag) + } + return nil +} diff --git a/pkg/yqlib/operator_divide_test.go b/pkg/yqlib/operator_divide_test.go new file mode 100644 index 00000000..d88106ed --- /dev/null +++ b/pkg/yqlib/operator_divide_test.go @@ -0,0 +1,121 @@ +package yqlib + +import ( + "testing" +) + +var divideOperatorScenarios = []expressionScenario{ + { + skipDoc: true, + document: `[{a: foo_bar, b: _}, {a: 4, b: 2}]`, + expression: ".[] | .a / .b", + expected: []string{ + "D0, P[0 a], (!!seq)::- foo\n- bar\n", + "D0, P[1 a], (!!float)::2\n", + }, + }, + { + skipDoc: true, + document: `{}`, + expression: "(.a / .b) as $x | .", + expected: []string{ + "D0, P[], (doc)::{}\n", + }, + }, + { + description: "String split", + document: `{a: cat_meow, b: _}`, + expression: `.c = .a / .b`, + expected: []string{ + "D0, P[], (doc)::{a: cat_meow, b: _, c: [cat, meow]}\n", + }, + }, + { + description: "Number division", + subdescription: "The result during divison is calculated as a float", + document: `{a: 12, b: 2.5}`, + expression: `.a = .a / .b`, + expected: []string{ + "D0, P[], (doc)::{a: 4.8, b: 2.5}\n", + }, + }, + { + skipDoc: true, + description: "Custom types: that are really strings", + document: "a: !horse cat_meow\nb: !goat _", + expression: `.a = .a / .b`, + expected: []string{ + "D0, P[], (doc)::a: !horse\n - cat\n - meow\nb: !goat _\n", + }, + }, + { + skipDoc: true, + description: "Custom types: that are really numbers", + document: "a: !horse 1.2\nb: !goat 2.3", + expression: `.a = .a / .b`, + expected: []string{ + "D0, P[], (doc)::a: !horse 0.5217391304347826\nb: !goat 2.3\n", + }, + }, + { + skipDoc: true, + document: "a: 2\nb: !goat 2.3", + expression: `.a = .a / .b`, + expected: []string{ + "D0, P[], (doc)::a: 0.8695652173913044\nb: !goat 2.3\n", + }, + }, + { + skipDoc: true, + description: "Custom types: that are really ints", + document: "a: !horse 2\nb: !goat 3", + expression: `.a += .b`, + expected: []string{ + "D0, P[], (doc)::a: !horse 5\nb: !goat 3\n", + }, + }, + { + skipDoc: true, + description: "Keep anchors", + document: "a: &horse [1]", + expression: `.a[1] = .a[0] / 2`, + expected: []string{ + "D0, P[], (doc)::a: &horse [1, 0.5]\n", + }, + }, + { + skipDoc: true, + description: "Divide int by string", + document: "a: 123\nb: '2'", + expression: `.a / .b`, + expectedError: "!!int cannot be divided by !!str", + }, + { + skipDoc: true, + description: "Divide string by int", + document: "a: 2\nb: '123'", + expression: `.b / .a`, + expectedError: "!!str cannot be divided by !!int", + }, + { + skipDoc: true, + description: "Divide map by int", + document: "a: {\"a\":1}\nb: 2", + expression: `.a / .b`, + expectedError: "!!map (a) cannot be divided by !!int (b)", + }, + { + skipDoc: true, + description: "Divide array by str", + document: "a: [1,2]\nb: '2'", + expression: `.a / .b`, + expectedError: "!!seq (a) cannot be divided by !!str (b)", + }, +} + +func TestDivideOperatorScenarios(t *testing.T) { + for _, tt := range divideOperatorScenarios { + testScenario(t, &tt) + } + documentOperatorScenarios(t, "divide", divideOperatorScenarios) +}