refactor(treesitter): get_{nodes,captures}_at_{position,cursor}

2 files changed, 89 insertions, 15 deletions

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 8d5e494601..afbece1088 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -370,6 +370,15 @@ attribute: >
 ==============================================================================
 Lua module: vim.treesitter                               *lua-treesitter-core*
 
+get_captures_at_cursor({winnr})                     *get_captures_at_cursor()*
+    Gets a list of captures under the cursor
+
+    Parameters: ~
+        {winnr}  (number|nil) Window handle or 0 for current window (default)
+
+    Return: ~
+        (table) Named node under the cursor
+
                                                   *get_captures_at_position()*
 get_captures_at_position({bufnr}, {row}, {col})
     Gets a list of captures for a given cursor position
@@ -380,7 +389,31 @@ get_captures_at_position({bufnr}, {row}, {col})
         {col}    (number) Position column
 
     Return: ~
-        (table) A table of captures
+        (table) Table of captures
+
+get_node_at_cursor({winnr})                             *get_node_at_cursor()*
+    Gets the smallest named node under the cursor
+
+    Parameters: ~
+        {winnr}  (number|nil) Window handle or 0 for current window (default)
+
+    Return: ~
+        (string) Named node under the cursor
+
+                                                      *get_node_at_position()*
+get_node_at_position({bufnr}, {row}, {col}, {opts})
+    Gets the smallest named node at position
+
+    Parameters: ~
+        {bufnr}  (number) Buffer number (0 for current buffer)
+        {row}    (number) Position row
+        {col}    (number) Position column
+        {opts}   (table) Optional keyword arguments:
+                 • ignore_injections boolean Ignore injected languages
+                   (default true)
+
+    Return: ~
+        (table) Named node under the cursor
 
 get_node_range({node_or_range})                             *get_node_range()*
     Get the node's range or unpack a range table
@@ -389,7 +422,7 @@ get_node_range({node_or_range})                             *get_node_range()*
         {node_or_range}  (table)
 
     Return: ~
-        start_row, start_col, end_row, end_col
+        (table) start_row, start_col, end_row, end_col
 
 get_parser({bufnr}, {lang}, {opts})                             *get_parser()*
     Gets the parser for this bufnr / ft combination.
@@ -398,7 +431,7 @@ get_parser({bufnr}, {lang}, {opts})                             *get_parser()*
     callback
 
     Parameters: ~
-        {bufnr}  (number|nil) Buffer the parser should be tied to: (default
+        {bufnr}  (number|nil) Buffer the parser should be tied to (default:
                  current buffer)
         {lang}   (string) |nil Filetype of this parser (default: buffer
                  filetype)
diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index d71f8611c1..8a540b9012 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -147,7 +147,7 @@ end
 ---
 ---@param node_or_range table
 ---
----@returns start_row, start_col, end_row, end_col
+---@returns table start_row, start_col, end_row, end_col
 function M.get_node_range(node_or_range)
   if type(node_or_range) == 'table' then
     return unpack(node_or_range)
@@ -198,7 +198,11 @@ end
 ---@param row number Position row
 ---@param col number Position column
 ---
----@returns (table) A table of captures
+---@param bufnr number Buffer number (0 for current buffer)
+---@param row number Position row
+---@param col number Position column
+---
+---@returns (table) Table of captures
 function M.get_captures_at_position(bufnr, row, col)
   if bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
@@ -245,25 +249,62 @@ function M.get_captures_at_position(bufnr, row, col)
   return matches
 end
 
---- Gets the smallest named node under the cursor
+---Gets a list of captures under the cursor
 ---
----@param winnr number Window handle or 0 for current window
----@param opts table Options table
----@param opts.ignore_injections boolean (default true) Ignore injected languages.
+---@param winnr number|nil Window handle or 0 for current window (default)
 ---
----@returns (table) The named node under the cursor
-function M.get_node_at_cursor(winnr, opts)
+---@returns (table) Named node under the cursor
+function M.get_captures_at_cursor(winnr)
   winnr = winnr or 0
+  local bufnr = a.nvim_win_get_buf(winnr)
   local cursor = a.nvim_win_get_cursor(winnr)
-  local ts_cursor_range = { cursor[1] - 1, cursor[2], cursor[1] - 1, cursor[2] }
 
-  local buf = a.nvim_win_get_buf(winnr)
-  local root_lang_tree = M.get_parser(buf)
+  local data = M.get_captures_at_position(bufnr, cursor[1] - 1, cursor[2])
+
+  local captures = {}
+
+  for _, capture in ipairs(data) do
+    table.insert(captures, capture.capture)
+  end
+
+  return captures
+end
+
+--- Gets the smallest named node at position
+---
+---@param bufnr number Buffer number (0 for current buffer)
+---@param row number Position row
+---@param col number Position column
+---@param opts table Optional keyword arguments:
+---             - ignore_injections boolean Ignore injected languages (default true)
+---
+---@returns (table) Named node under the cursor
+function M.get_node_at_position(bufnr, row, col, opts)
+  if bufnr == 0 then
+    bufnr = a.nvim_get_current_buf()
+  end
+  local ts_range = { row, col, row, col }
+
+  local root_lang_tree = M.get_parser(bufnr)
   if not root_lang_tree then
     return
   end
 
-  return root_lang_tree:named_node_for_range(ts_cursor_range, opts)
+  return root_lang_tree:named_node_for_range(ts_range, opts)
+end
+
+--- Gets the smallest named node under the cursor
+---
+---@param winnr number|nil Window handle or 0 for current window (default)
+---
+---@returns (string) Named node under the cursor
+function M.get_node_at_cursor(winnr)
+  winnr = winnr or 0
+  local bufnr = a.nvim_win_get_buf(winnr)
+  local cursor = a.nvim_win_get_cursor(winnr)
+
+  return M.get_node_at_position(bufnr, cursor[1] - 1, cursor[2], { ignore_injections = false })
+    :type()
 end
 
 --- Start treesitter highlighting for a buffer