feat(treesitter): provide select() #39069

Problem: No public method for treesitter incremental selection.

Solution: Add `vim.treesitter.select()`.

runtime/doc/news.txt

@@ -216,6 +216,7 @@ TREESITTER
 • |v_]N| |v_[N| expand selection to sibling treesitter node.
 • |treesitter-highlight-conceal| can be removed by adding a `@noconceal`
   capture.
+• |vim.treesitter.select()| starts or adjusts a visual selection at cursor, based on tree nodes.
 
 TUI
 

runtime/doc/treesitter.txt

@@ -1160,6 +1160,16 @@ node_contains({node}, {range})                *vim.treesitter.node_contains()*
     Return: ~
         (`boolean`) True if the {node} contains the {range}
 
+select({direction}, {opts})                          *vim.treesitter.select()*
+    Starts or adjusts a |Visual| selection at cursor, based on tree nodes.
+
+    Parameters: ~
+      • {direction}  (`'parent'|'child'|'next'|'prev'|'extend_next'|'extend_prev'`)
+                     Direction to select towards
+      • {opts}       (`table?`) A table with the following fields:
+                     • {count} (`integer?`) Number of selections to make in
+                       the given direction (default 1)
+
 start({buf}, {lang})                                  *vim.treesitter.start()*
     Starts treesitter highlighting for a buffer
 

runtime/lua/vim/treesitter.lua

@@ -12,6 +12,7 @@ local M = vim._defer_require('vim.treesitter', {
   language = ..., --- @module 'vim.treesitter.language'
   languagetree = ..., --- @module 'vim.treesitter.languagetree'
   query = ..., --- @module 'vim.treesitter.query'
+  _select = ..., --- @module 'vim.treesitter._select'
 })
 
 local LanguageTree = M.languagetree
@@ -508,4 +509,42 @@ function M.foldexpr(lnum)
   return M._fold.foldexpr(lnum)
 end
 
+--- @class vim.treesitter.select.Opts
+--- @inlinedoc
+---
+--- Number of selections to make in the given direction (default 1)
+--- @field count integer?
+
+--- Starts or adjusts a |Visual| selection at cursor, based on tree nodes.
+---@param direction 'parent'|'child'|'next'|'prev'|'extend_next'|'extend_prev' Direction to select
+---                                                                            towards
+---@param opts vim.treesitter.select.Opts?
+function M.select(direction, opts)
+  vim.validate('direction', direction, 'string')
+  vim.validate('opts', opts, 'table', true)
+  opts = opts or {}
+  if opts.count then
+    vim.validate('count', opts.count, 'number')
+  end
+  local count = opts.count or 1
+
+  if direction == 'parent' then
+    return M._select.select_parent(count)
+  elseif direction == 'child' then
+    return M._select.select_child(count)
+  elseif direction == 'next' then
+    return M._select.select_next(count)
+  elseif direction == 'prev' then
+    return M._select.select_prev(count)
+  elseif direction == 'extend_next' then
+    return M._select.select_grow_next(count)
+  elseif direction == 'extend_prev' then
+    return M._select.select_grow_prev(count)
+  else
+    vim.validate('direction', direction, function()
+      return false, 'Invalid direction'
+    end)
+  end
+end
+
 return M

test/functional/treesitter/select_spec.lua

@@ -23,14 +23,26 @@ local function set_filetype(ft)
   api.nvim_set_option_value('filetype', ft, { buf = 0 })
 end
 
-local function treeselect(cmd_, ...)
+local function treeselect(cmd_, count_)
   if cmd_ == 'select_node' then
-    cmd_ = 'select_child'
+    cmd_ = 'child'
+  elseif cmd_ == 'select_child' then
+    cmd_ = 'child'
+  elseif cmd_ == 'select_parent' then
+    cmd_ = 'parent'
+  elseif cmd_ == 'select_next' then
+    cmd_ = 'next'
+  elseif cmd_ == 'select_prev' then
+    cmd_ = 'prev'
+  elseif cmd_ == 'select_grow_next' then
+    cmd_ = 'extend_next'
+  elseif cmd_ == 'select_grow_prev' then
+    cmd_ = 'extend_prev'
   end
 
-  exec_lua(function(cmd, ...)
-    require 'vim.treesitter._select'[cmd](...)
-  end, cmd_, ...)
+  exec_lua(function(cmd, count)
+    vim.treesitter.select(cmd, { count = count })
+  end, cmd_, count_)
 end
 
 describe('treesitter incremental-selection', function()