Merge pull request #20093 from neovim/upstream-ts-util

* feat(treesitter): upstream get_{node,captures} utility functions
* refactor(treesitter): get_{nodes,captures}_at_{position,cursor}

1 files changed, 65 insertions, 3 deletions

diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 69faea7edc..d93c485dfe 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -78,7 +78,7 @@ end
 --- If needed this will create the parser.
 --- Unconditionally attach the provided callback
 ---
----@param bufnr number|nil Buffer the parser should be tied to: (default current buffer)
+---@param bufnr number|nil Buffer the parser should be tied to (default: current buffer)
 ---@param lang string |nil Filetype of this parser (default: buffer filetype)
 ---@param opts table|nil Options to pass to the created language tree
 ---
@@ -144,7 +144,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)
@@ -195,7 +195,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()
@@ -242,6 +246,64 @@ function M.get_captures_at_position(bufnr, row, col)
   return matches
 end
 
+---Gets a list of captures under the cursor
+---
+---@param winnr number|nil Window handle or 0 for current window (default)
+---
+---@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 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_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
 ---
 --- Can be used in an ftplugin or FileType autocommand