diff --git a/builtin_metadata.json b/builtin_metadata.json index 3e658c012b..2558ff8ebd 100644 --- a/builtin_metadata.json +++ b/builtin_metadata.json @@ -185,6 +185,7 @@ "strings.render_template", "strings.replace_n", "strings.reverse", + "strings.split_n", "substring", "trim", "trim_left", @@ -24640,6 +24641,36 @@ }, "wasm": true }, + "strings.split_n": { + "args": [ + { + "description": "string that is split", + "name": "x", + "type": "string" + }, + { + "description": "delimiter used for splitting", + "name": "delimiter", + "type": "string" + }, + { + "description": "number of parts to return; positive selects from the left, negative from the right, zero returns an empty array", + "name": "n", + "type": "number" + } + ], + "available": [ + "edge" + ], + "description": "Returns an array of at most `n` parts of `x` split on `delimiter`. If `n` is positive, returns the first `n` parts. If `n` is negative, returns the last `abs(n)` parts. If `n` is zero, returns an empty array. If `abs(n)` exceeds the number of parts, all parts are returned.", + "introduced": "edge", + "result": { + "description": "split parts", + "name": "ys", + "type": "array[string]" + }, + "wasm": false + }, "substring": { "args": [ { diff --git a/capabilities.json b/capabilities.json index c58359321b..1b02dd4ace 100644 --- a/capabilities.json +++ b/capabilities.json @@ -4177,6 +4177,29 @@ "type": "function" } }, + { + "name": "strings.split_n", + "decl": { + "args": [ + { + "type": "string" + }, + { + "type": "string" + }, + { + "type": "number" + } + ], + "result": { + "dynamic": { + "type": "string" + }, + "type": "array" + }, + "type": "function" + } + }, { "name": "substring", "decl": { diff --git a/v1/ast/builtins.go b/v1/ast/builtins.go index 56661f7dce..17ed06035d 100644 --- a/v1/ast/builtins.go +++ b/v1/ast/builtins.go @@ -141,6 +141,7 @@ var DefaultBuiltins = [...]*Builtin{ StartsWith, EndsWith, Split, + SplitN, Replace, ReplaceN, Trim, @@ -1281,6 +1282,21 @@ var Split = &Builtin{ CanSkipBctx: true, } +var SplitN = &Builtin{ + Name: "strings.split_n", + Description: "Returns an array of at most `n` parts of `x` split on `delimiter`. If `n` is positive, returns the first `n` parts. If `n` is negative, returns the last `abs(n)` parts. If `n` is zero, returns an empty array. If `abs(n)` exceeds the number of parts, all parts are returned.", + Decl: types.NewFunction( + types.Args( + types.Named("x", types.S).Description("string that is split"), + types.Named("delimiter", types.S).Description("delimiter used for splitting"), + types.Named("n", types.N).Description("number of parts to return; positive selects from the left, negative from the right, zero returns an empty array"), + ), + types.Named("ys", types.NewArray(nil, types.S)).Description("split parts"), + ), + Categories: stringsCat, + CanSkipBctx: true, +} + var Replace = &Builtin{ Name: "replace", Description: "Replace replaces all instances of a sub-string.", diff --git a/v1/test/cases/testdata/v1/strings/test-splitn.yaml b/v1/test/cases/testdata/v1/strings/test-splitn.yaml new file mode 100644 index 0000000000..353c702d46 --- /dev/null +++ b/v1/test/cases/testdata/v1/strings/test-splitn.yaml @@ -0,0 +1,194 @@ +--- +cases: + - note: strings/split_n/positive_basic + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", 2) + } + want_result: + - x: ["a", "b"] + - note: strings/split_n/positive_one + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", 1) + } + want_result: + - x: ["a"] + - note: strings/split_n/positive_exact + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", 4) + } + want_result: + - x: ["a", "b", "c", "d"] + - note: strings/split_n/positive_more_than_parts + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", 10) + } + want_result: + - x: ["a", "b", "c", "d"] + - note: strings/split_n/zero + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", 0) + } + want_result: + - x: [] + - note: strings/split_n/negative_basic + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", -2) + } + want_result: + - x: ["c", "d"] + - note: strings/split_n/negative_one + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", -1) + } + want_result: + - x: ["d"] + - note: strings/split_n/negative_exact + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", -4) + } + want_result: + - x: ["a", "b", "c", "d"] + - note: strings/split_n/negative_more_than_parts + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c.d", ".", -10) + } + want_result: + - x: ["a", "b", "c", "d"] + - note: strings/split_n/no_delimiter_found_positive + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("abc", ".", 2) + } + want_result: + - x: ["abc"] + - note: strings/split_n/no_delimiter_found_negative + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("abc", ".", -2) + } + want_result: + - x: ["abc"] + - note: strings/split_n/empty_string + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("", ".", 2) + } + want_result: + - x: [""] + - note: strings/split_n/empty_delimiter_positive + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("abc", "", 2) + } + want_result: + - x: ["a", "b"] + - note: strings/split_n/empty_delimiter_negative + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("abc", "", -2) + } + want_result: + - x: ["b", "c"] + - note: strings/split_n/non_integer_n + query: data.test.p = x + modules: + - | + package test + + p := result if { + result := strings.split_n("a.b.c", ".", 1.5) + } + strict_error: true + want_error_code: eval_type_error + want_error: "strings.split_n: operand 3 must be integer number but got floating-point number" + - note: strings/split_n/usecase_positive + query: data.test.p = x + modules: + - | + package test + + p := [name, email] if { + [name, email] := strings.split_n(input.user, ";", 2) + } + input: + user: "alice;alice@example.com;admin;active" + want_result: + - x: ["alice", "alice@example.com"] + - note: strings/split_n/usecase_negative + query: data.test.p = x + modules: + - | + package test + + p := [name, email] if { + [name, email] := strings.split_n(input.user, ";", -2) + } + input: + user: "admin;active;alice;alice@example.com" + want_result: + - x: ["alice", "alice@example.com"] diff --git a/v1/topdown/strings.go b/v1/topdown/strings.go index 99885622f5..4ffd307c78 100644 --- a/v1/topdown/strings.go +++ b/v1/topdown/strings.go @@ -557,6 +557,55 @@ func builtinSplit(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) e return iter(ast.ArrayTerm(util.SplitMap(text, delim, ast.InternedTerm)...)) } +func builtinSplitN(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { + s, err := builtins.StringOperand(operands[0].Value, 1) + if err != nil { + return err + } + + d, err := builtins.StringOperand(operands[1].Value, 2) + if err != nil { + return err + } + + n, err := builtins.IntOperand(operands[2].Value, 3) + if err != nil { + return err + } + + text, delim := string(s), string(d) + + var result []*ast.Term + if n >= 0 { + // n+1 may overflow for very large n; a negative limit means no limit. + limit := n + 1 + if limit < 0 { + limit = -1 + } + parts := strings.SplitN(text, delim, limit) + end := n + if end > len(parts) { + end = len(parts) + } + result = make([]*ast.Term, end) + for i := range result { + result[i] = ast.InternedTerm(parts[i]) + } + } else { + parts := strings.Split(text, delim) + start := len(parts) + n + if start < 0 { + start = 0 + } + result = make([]*ast.Term, len(parts)-start) + for i, p := range parts[start:] { + result[i] = ast.InternedTerm(p) + } + } + + return iter(ast.ArrayTerm(result...)) +} + func builtinReplace(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { s, err := builtins.StringOperand(operands[0].Value, 1) if err != nil { @@ -831,6 +880,7 @@ func init() { RegisterBuiltinFunc(ast.Upper.Name, builtinUpper) RegisterBuiltinFunc(ast.Lower.Name, builtinLower) RegisterBuiltinFunc(ast.Split.Name, builtinSplit) + RegisterBuiltinFunc(ast.SplitN.Name, builtinSplitN) RegisterBuiltinFunc(ast.Replace.Name, builtinReplace) RegisterBuiltinFunc(ast.ReplaceN.Name, builtinReplaceN) RegisterBuiltinFunc(ast.Trim.Name, builtinTrim)