diff options
| -rw-r--r-- | runtime/doc/treesitter.txt | 5 | ||||
| -rw-r--r-- | runtime/lua/vim/treesitter/query.lua | 20 | ||||
| -rw-r--r-- | test/functional/treesitter/parser_spec.lua | 61 |
3 files changed, 84 insertions, 2 deletions
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt index 1f4b5d3097..39522898f9 100644 --- a/runtime/doc/treesitter.txt +++ b/runtime/doc/treesitter.txt @@ -212,6 +212,11 @@ Here is a list of built-in predicates : ((identifier) @foo (#contains? @foo "foo")) ((identifier) @foo-bar (#contains @foo-bar "foo" "bar")) < + `any-of?` *ts-predicate-any-of?* + Will check if the text is the same as any of the following + arguments : > + ((identifier) @foo (#any-of? @foo "foo" "bar")) +< *lua-treesitter-not-predicate* Each predicate has a `not-` prefixed predicate that is just the negation of the predicate. diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua index c0140f9186..b81eb18945 100644 --- a/runtime/lua/vim/treesitter/query.lua +++ b/runtime/lua/vim/treesitter/query.lua @@ -260,7 +260,25 @@ local predicate_handlers = { end return false - end + end, + + ["any-of?"] = function(match, _, source, predicate) + local node = match[predicate[2]] + local node_text = M.get_node_text(node, source) + + -- Since 'predicate' will not be used by callers of this function, use it + -- to store a string set built from the list of words to check against. + local string_set = predicate["string_set"] + if not string_set then + string_set = {} + for i=3,#predicate do + string_set[predicate[i]] = true + end + predicate["string_set"] = string_set + end + + return string_set[node_text] + end, } -- As we provide lua-match? also expose vim-match? diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua index 1017913709..d2f9148e8f 100644 --- a/test/functional/treesitter/parser_spec.lua +++ b/test/functional/treesitter/parser_spec.lua @@ -270,6 +270,65 @@ void ui_refresh(void) }, res) end) + it('supports builtin query predicate any-of?', function() + if pending_c_parser(pending) then return end + + insert([[ + #include <stdio.h> + + int main(void) { + int i; + for(i=1; i<=100; i++) { + if(((i%3)||(i%5))== 0) + printf("number= %d FizzBuzz\n", i); + else if((i%3)==0) + printf("number= %d Fizz\n", i); + else if((i%5)==0) + printf("number= %d Buzz\n", i); + else + printf("number= %d\n",i); + } + return 0; + } + ]]) + exec_lua([[ + function get_query_result(query_text) + cquery = vim.treesitter.parse_query("c", query_text) + parser = vim.treesitter.get_parser(0, "c") + tree = parser:parse()[1] + res = {} + for cid, node in cquery:iter_captures(tree:root(), 0) do + -- can't transmit node over RPC. just check the name, range, and text + local text = vim.treesitter.get_node_text(node, 0) + local range = {node:range()} + table.insert(res, {cquery.captures[cid], node:type(), range, text}) + end + return res + end + ]]) + + local res0 = exec_lua([[return get_query_result(...)]], + [[((primitive_type) @c-keyword (#any-of? @c-keyword "int" "float"))]]) + eq({ + { "c-keyword", "primitive_type", { 2, 2, 2, 5 }, "int" }, + { "c-keyword", "primitive_type", { 3, 4, 3, 7 }, "int" }, + }, res0) + + local res1 = exec_lua([[return get_query_result(...)]], + [[ + ((string_literal) @fizzbuzz-strings (#any-of? @fizzbuzz-strings + "\"number= %d FizzBuzz\\n\"" + "\"number= %d Fizz\\n\"" + "\"number= %d Buzz\\n\"" + )) + ]]) + eq({ + { "fizzbuzz-strings", "string_literal", { 6, 15, 6, 38 }, "\"number= %d FizzBuzz\\n\""}, + { "fizzbuzz-strings", "string_literal", { 8, 15, 8, 34 }, "\"number= %d Fizz\\n\""}, + { "fizzbuzz-strings", "string_literal", { 10, 15, 10, 34 }, "\"number= %d Buzz\\n\""}, + }, res1) + end) + it('allow loading query with escaped quotes and capture them with `lua-match?` and `vim-match?`', function() if pending_c_parser(pending) then return end @@ -343,7 +402,7 @@ void ui_refresh(void) return list ]] - eq({ 'contains?', 'eq?', 'is-main?', 'lua-match?', 'match?', 'vim-match?' }, res_list) + eq({ 'any-of?', 'contains?', 'eq?', 'is-main?', 'lua-match?', 'match?', 'vim-match?' }, res_list) end) |