Range.from_tsquery_caps

lua/u/range.lua

@@ -228,6 +228,40 @@ function Range.from_motion(motion, opts)
   return captured_range
 end
 
+--- @param opts? { contains_cursor?: boolean }
+function Range.from_tsquery_caps(bufnr, query, opts)
+  if bufnr == nil or bufnr == 0 then bufnr = vim.api.nvim_get_current_buf() end
+  local lang = vim.treesitter.language.get_lang(vim.bo[bufnr].filetype)
+  if lang == nil then return end
+  local parser = vim.treesitter.get_parser(bufnr, lang)
+  if parser == nil then return end
+  local tree = parser:parse()[1]
+  if tree == nil then return end
+
+  opts = opts or { contains_cursor = true }
+  local cursor = Pos.from_pos '.'
+
+  local root = tree:root()
+  local root_end = root:end_()
+  local q = vim.treesitter.query.parse(lang, query)
+  --- @type table<string, u.Range[]>
+  local ranges = {}
+  for id, match, _meta in q:iter_captures(root, bufnr, 0, root_end) do
+    local start_row0, start_col0, stop_row0, stop_col0 = match:range()
+    local range = Range.new(
+      Pos.new(bufnr, start_row0 + 1, start_col0 + 1),
+      Pos.new(bufnr, stop_row0 + 1, stop_col0 + 1),
+      'v'
+    )
+    if not opts.contains_cursor or opts.contains_cursor and range:contains(cursor) then
+      local capture_name = q.captures[id]
+      if not ranges[capture_name] then ranges[capture_name] = {} end
+      table.insert(ranges[capture_name], range)
+    end
+  end
+  return ranges
+end
+
 --- Get range information from the currently selected visual text.
 --- Note: from within a command mapping or an opfunc, use other specialized
 --- utilities, such as: