diff --git a/ast/builtins.go b/ast/builtins.go index d0ab69a163..c5f70e7f5c 100644 --- a/ast/builtins.go +++ b/ast/builtins.go @@ -206,6 +206,8 @@ var Upper = v1.Upper var Split = v1.Split +var SplitN = v1.SplitN + var Replace = v1.Replace var ReplaceN = v1.ReplaceN diff --git a/builtin_metadata.json b/builtin_metadata.json index f176a0103e..c572716cde 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", @@ -22620,6 +22621,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": "positive: at most n substrings from left; negative: last |n| substrings from right; zero: nil", + "name": "n", + "type": "number" + } + ], + "available": [ + "edge" + ], + "description": "Split returns an array containing elements of the input string split on a delimiter. A positive n limits the number of substrings returned from the left. A negative n returns the last |n| substrings from the right. Zero returns nil.", + "introduced": "edge", + "result": { + "description": "split parts", + "name": "ys", + "type": "array[string]" + }, + "wasm": false + }, "substring": { "args": [ { diff --git a/capabilities.json b/capabilities.json index 9eb82c2968..5fd9f900ef 100644 --- a/capabilities.json +++ b/capabilities.json @@ -4129,6 +4129,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 7e30a8051c..7467274966 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, @@ -1279,6 +1280,21 @@ var Split = &Builtin{ CanSkipBctx: true, } +var SplitN = &Builtin{ + Name: "strings.split_n", + Description: "Split returns an array containing elements of the input string split on a delimiter. A positive n limits the number of substrings returned from the left. A negative n returns the last |n| substrings from the right. Zero returns nil.", + 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("positive: at most n substrings from left; negative: last |n| substrings from right; zero: nil"), + ), + 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..aafbb14b3a --- /dev/null +++ b/v1/test/cases/testdata/v1/strings/test-splitn.yaml @@ -0,0 +1,101 @@ +--- +cases: + - note: "strings/split_n: basic split with limit" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a.b.c.d", ".", 2) + want_result: + - x: ["a", "b.c.d"] + - note: "strings/split_n: limit greater than splits" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a.b", ".", 5) + want_result: + - x: ["a", "b"] + - note: "strings/split_n: negative limit takes last n from right" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a.b.c", ".", -1) + want_result: + - x: ["c"] + - note: "strings/split_n: negative limit takes last 2 from right" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a;b;c;d", ";", -2) + want_result: + - x: ["c", "d"] + - note: "strings/split_n: negative limit exceeds splits returns all" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a.b", ".", -5) + want_result: + - x: ["a", "b"] + - note: "strings/split_n: zero limit returns nil" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a.b.c", ".", 0) + want_result: + - x: [] + - note: "strings/split_n: limit of 1 returns whole string" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a.b.c", ".", 1) + want_result: + - x: ["a.b.c"] + - note: "strings/split_n: no delimiter found" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("abc", ".", 2) + want_result: + - x: ["abc"] + - note: "strings/split_n: empty string" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("", ".", 2) + want_result: + - x: [""] + - note: "strings/split_n: empty delimiter" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("abc", "", 2) + want_result: + - x: ["a", "bc"] + - note: "strings/split_n: multi-char delimiter" + query: data.test.p = x + modules: + - | + package test + + p := strings.split_n("a::b::c::d", "::", 3) + want_result: + - x: ["a", "b", "c::d"] diff --git a/v1/topdown/strings.go b/v1/topdown/strings.go index 27c841102b..00d05f7096 100644 --- a/v1/topdown/strings.go +++ b/v1/topdown/strings.go @@ -538,6 +538,44 @@ 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 elems []string + if n < 0 { + // Negative n: split all, then take the last |n| elements. + all := strings.Split(text, delim) + if -n >= len(all) { + elems = all + } else { + elems = all[len(all)+n:] + } + } else { + elems = strings.SplitN(text, delim, n) + } + + arr := make([]*ast.Term, len(elems)) + for i := range elems { + arr[i] = ast.InternedTerm(elems[i]) + } + return iter(ast.ArrayTerm(arr...)) +} + func builtinReplace(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { s, err := builtins.StringOperand(operands[0].Value, 1) if err != nil { @@ -801,6 +839,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)