diff options
-rw-r--r-- | src/nvim/lua/treesitter.c | 26 | ||||
-rw-r--r-- | test/functional/treesitter/node_spec.lua | 18 |
2 files changed, 44 insertions, 0 deletions
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c index 954c3410d3..2ef4d6052c 100644 --- a/src/nvim/lua/treesitter.c +++ b/src/nvim/lua/treesitter.c @@ -88,6 +88,7 @@ static struct luaL_Reg node_meta[] = { { "prev_sibling", node_prev_sibling }, { "next_named_sibling", node_next_named_sibling }, { "prev_named_sibling", node_prev_named_sibling }, + { "named_children", node_named_children }, { NULL, NULL } }; @@ -1062,6 +1063,31 @@ static int node_prev_named_sibling(lua_State *L) return 1; } +static int node_named_children(lua_State *L) +{ + TSNode source; + if (!node_check(L, 1, &source)) { + return 0; + } + TSTreeCursor cursor = ts_tree_cursor_new(source); + + lua_newtable(L); + int curr_index = 0; + + if (ts_tree_cursor_goto_first_child(&cursor)) { + do { + TSNode node = ts_tree_cursor_current_node(&cursor); + if (ts_node_is_named(node)) { + push_node(L, node, 1); + lua_rawseti(L, -2, ++curr_index); + } + } while (ts_tree_cursor_goto_next_sibling(&cursor)); + } + + ts_tree_cursor_delete(&cursor); + return 1; +} + /// assumes the match table being on top of the stack static void set_match(lua_State *L, TSQueryMatch *match, int nodeidx) { diff --git a/test/functional/treesitter/node_spec.lua b/test/functional/treesitter/node_spec.lua index 21c287644e..b9ad4925c9 100644 --- a/test/functional/treesitter/node_spec.lua +++ b/test/functional/treesitter/node_spec.lua @@ -59,4 +59,22 @@ describe('treesitter node API', function() exec_lua 'node = node:prev_named_sibling()' eq('int x', lua_eval('node_text(node)')) end) + + it('can retrieve the children of a node', function() + insert([[ + int main() { + int x = 3; + }]]) + + local len = exec_lua([[ + tree = vim.treesitter.get_parser(0, "c"):parse()[1] + node = tree:root():child(0) + children = node:named_children() + + return #children + ]]) + + eq(3, len) + eq('<node compound_statement>', lua_eval('tostring(children[3])')) + end) end) |