feat(treesitter): allow disabling captures and patterns on TSQuery (#32790)

Problem: Cannot disable individual captures and patterns in treesitter queries.

Solution: 
* Expose the corresponding tree-sitter API functions for `TSQuery` object. 
* Add documentation for `TSQuery`.
* Return the pattern ID from `get_captures_at_pos()` (and hence `:Inspect!`).

runtime/doc/news.txt

@@ -429,6 +429,10 @@ TREESITTER
   code block fence lines vertically.
 • |vim.treesitter.language.inspect()| shows additional information, including
   parser version for ABI 15 parsers.
+• |TSQuery:disable_pattern()| and |TSQuery:disable_capture()| to turn off
+  a specific pattern or capture in a query.
+• |vim.treesitter.get_captures_at_pos()| returns the `pattern_id` of the
+  pattern used to match each capture.
 
 TUI
 

runtime/doc/treesitter.txt

@@ -1240,6 +1240,26 @@ This Lua |treesitter-query| interface allows you to create queries and use
 them to parse text. See |vim.treesitter.query.parse()| for a working example.
 
 
+*vim.treesitter.Query*
+    Parsed query, see |vim.treesitter.query.parse()|
+
+    Fields: ~
+      • {lang}                     (`string`) parser language name
+      • {captures}                 (`string[]`) list of (unique) capture names
+                                   defined in query
+      • {info}                     (`vim.treesitter.QueryInfo`) query context
+                                   (e.g. captures, predicates, directives)
+      • {has_conceal_line}         (`boolean`) whether the query sets
+                                   conceal_lines metadata
+      • {has_combined_injections}  (`boolean`) whether the query contains
+                                   combined injections
+      • {query}                    (`TSQuery`) userdata query object
+      • {iter_captures}            (`fun(self: vim.treesitter.Query, node: TSNode, source: integer|string, start: integer?, stop: integer?): fun(end_line: integer?): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch, TSTree`)
+                                   See |Query:iter_captures()|.
+      • {iter_matches}             (`fun(self: vim.treesitter.Query, node: TSNode, source: integer|string, start: integer?, stop: integer?, opts: table?): fun(): integer, table<integer, TSNode[]>, vim.treesitter.query.TSMetadata, TSTree`)
+                                   See |Query:iter_matches()|.
+
+
                                         *vim.treesitter.query.add_directive()*
 add_directive({name}, {handler}, {opts})
     Adds a new directive to be used in queries
@@ -1307,7 +1327,7 @@ get({lang}, {query_name})                         *vim.treesitter.query.get()*
 
     Return: ~
         (`vim.treesitter.Query?`) Parsed query. `nil` if no query files are
-        found.
+        found. See |vim.treesitter.Query|.
 
                                             *vim.treesitter.query.get_files()*
 get_files({lang}, {query_name}, {is_included})
@@ -1373,10 +1393,12 @@ parse({lang}, {query})                          *vim.treesitter.query.parse()*
     Parses a {query} string and returns a `Query` object
     (|lua-treesitter-query|), which can be used to search the tree for the
     query patterns (via |Query:iter_captures()|, |Query:iter_matches()|), or
-    inspect the query via these fields:
+    inspect/modify the query via these fields:
     • `captures`: a list of unique capture names defined in the query (alias:
       `info.captures`).
     • `info.patterns`: information about predicates.
+    • `query`: the underlying |TSQuery| which can be used to disable patterns
+      or captures.
 
     Example: >lua
         local query = vim.treesitter.query.parse('vimdoc', [[
@@ -1396,7 +1418,7 @@ parse({lang}, {query})                          *vim.treesitter.query.parse()*
       • {query}  (`string`) Query text, in s-expr syntax
 
     Return: ~
-        (`vim.treesitter.Query`) Parsed query
+        (`vim.treesitter.Query`) Parsed query . See |vim.treesitter.Query|.
 
     See also: ~
       • |vim.treesitter.query.get()|
@@ -1514,6 +1536,56 @@ set({lang}, {query_name}, {text})                 *vim.treesitter.query.set()*
       • {text}        (`string`) Query text (unparsed).
 
 
+
+
+*TSQuery*
+    Extends: |userdata|
+
+    Reference to an object held by the treesitter library that is used as a
+    component of the |vim.treesitter.Query| for language feature support. See
+    |treesitter-query| for more about queries or
+    |vim.treesitter.query.parse()| for an example of how to obtain a query
+    object.
+
+    Fields: ~
+      • {disable_capture}  (`fun(self: TSQuery, capture_name: string)`) See
+                           |TSQuery:disable_capture()|.
+      • {disable_pattern}  (`fun(self: TSQuery, pattern_index: integer)`) See
+                           |TSQuery:disable_pattern()|.
+
+
+TSQuery:disable_capture({capture_name})            *TSQuery:disable_capture()*
+    Disable a specific capture in this query; once disabled the capture cannot
+    be re-enabled. {capture_name} should not include a leading "@".
+
+    Example: To disable the `@variable.parameter` capture from the vimdoc
+    highlights query: >lua
+        local query = vim.treesitter.query.get('vimdoc', 'highlights')
+        query.query:disable_capture("variable.parameter")
+        vim.treesitter.get_parser():parse()
+<
+
+    Parameters: ~
+      • {capture_name}  (`string`)
+
+TSQuery:disable_pattern({pattern_index})           *TSQuery:disable_pattern()*
+    Disable a specific pattern in this query; once disabled the pattern cannot
+    be re-enabled. The {pattern_index} for a particular match can be obtained
+    with |:Inspect!|, or by reading the source of the query (i.e. from
+    |vim.treesitter.query.get_files()|).
+
+    Example: To disable `|` links in vimdoc but keep other `@markup.link`s
+    highlighted: >lua
+        local link_pattern = 9 -- from :Inspect!
+        local query = vim.treesitter.query.get('vimdoc', 'highlights')
+        query.query:disable_pattern(link_pattern)
+        local tree = vim.treesitter.get_parser():parse()[1]
+<
+
+    Parameters: ~
+      • {pattern_index}  (`integer`)
+
+
 ==============================================================================
 Lua module: vim.treesitter.languagetree              *treesitter-languagetree*
 

runtime/lua/vim/treesitter.lua

@@ -288,15 +288,19 @@ function M.get_captures_at_pos(bufnr, row, col)
 
     local iter = q:query():iter_captures(root, buf_highlighter.bufnr, row, row + 1)
 
-    for id, node, metadata in iter do
+    for id, node, metadata, match in iter do
       if M.is_in_node_range(node, row, col) then
         ---@diagnostic disable-next-line: invisible
         local capture = q._query.captures[id] -- name of the capture in the query
         if capture ~= nil then
-          table.insert(
-            matches,
-            { capture = capture, metadata = metadata, lang = tree:lang(), id = id }
-          )
+          local _, pattern_id = match:info()
+          table.insert(matches, {
+            capture = capture,
+            metadata = metadata,
+            lang = tree:lang(),
+            id = id,
+            pattern_id = pattern_id,
+          })
         end
       end
     end

runtime/lua/vim/treesitter/_meta/misc.lua

@@ -14,9 +14,6 @@ error('Cannot require a meta file')
 ---@field _set_logger fun(self: TSParser, lex: boolean, parse: boolean, cb: TSLoggerCallback)
 ---@field _logger fun(self: TSParser): TSLoggerCallback
 
----@class TSQuery: userdata
----@field inspect fun(self: TSQuery): TSQueryInfo
-
 ---@class (exact) TSQueryInfo
 ---@field captures string[]
 ---@field patterns table<integer, (integer|string)[][]>

runtime/lua/vim/treesitter/_meta/tsquery.lua

@@ -0,0 +1,45 @@
+---@meta
+-- luacheck: no unused args
+error('Cannot require a meta file')
+
+-- This could be documented as a module @brief like tsnode/tstree, but without
+-- its own section header documenting it as a class ensures it still gets a helptag.
+
+--- Reference to an object held by the treesitter library that is used as a
+--- component of the |vim.treesitter.Query| for language feature support.
+--- See |treesitter-query| for more about queries or |vim.treesitter.query.parse()|
+--- for an example of how to obtain a query object.
+---
+---@class TSQuery: userdata
+local TSQuery = {} -- luacheck: no unused
+
+--- Get information about the query's patterns and captures.
+---@nodoc
+---@return TSQueryInfo
+function TSQuery:inspect() end
+
+--- Disable a specific capture in this query; once disabled the capture cannot be re-enabled.
+--- {capture_name} should not include a leading "@".
+---
+--- Example: To disable the `@variable.parameter` capture from the vimdoc highlights query:
+--- ```lua
+--- local query = vim.treesitter.query.get('vimdoc', 'highlights')
+--- query.query:disable_capture("variable.parameter")
+--- vim.treesitter.get_parser():parse()
+--- ```
+---@param capture_name string
+function TSQuery:disable_capture(capture_name) end
+
+--- Disable a specific pattern in this query; once disabled the pattern cannot be re-enabled.
+--- The {pattern_index} for a particular match can be obtained with |:Inspect!|, or by reading
+--- the source of the query (i.e. from |vim.treesitter.query.get_files()|).
+---
+--- Example: To disable `|` links in vimdoc but keep other `@markup.link`s highlighted:
+--- ```lua
+--- local link_pattern = 9 -- from :Inspect!
+--- local query = vim.treesitter.query.get('vimdoc', 'highlights')
+--- query.query:disable_pattern(link_pattern)
+--- local tree = vim.treesitter.get_parser():parse()[1]
+--- ```
+---@param pattern_index integer
+function TSQuery:disable_pattern(pattern_index) end

runtime/lua/vim/treesitter/query.lua

@@ -10,7 +10,6 @@ local EXTENDS_FORMAT = '^;+%s*extends%s*$'
 
 local M = {}
 
----@nodoc
 ---Parsed query, see |vim.treesitter.query.parse()|
 ---
 ---@class vim.treesitter.Query
@@ -344,9 +343,10 @@ api.nvim_create_autocmd('OptionSet', {
 
 --- Parses a {query} string and returns a `Query` object (|lua-treesitter-query|), which can be used
 --- to search the tree for the query patterns (via |Query:iter_captures()|, |Query:iter_matches()|),
---- or inspect the query via these fields:
+--- or inspect/modify the query via these fields:
 ---   - `captures`: a list of unique capture names defined in the query (alias: `info.captures`).
 ---   - `info.patterns`: information about predicates.
+---   - `query`: the underlying |TSQuery| which can be used to disable patterns or captures.
 ---
 --- Example:
 --- ```lua

src/gen/gen_vimdoc.lua

@@ -328,10 +328,12 @@ local config = {
       'treesitter.lua',
       'language.lua',
       'query.lua',
+      'tsquery.lua',
       'highlighter.lua',
       'languagetree.lua',
       'dev.lua',
     },
+    append_only = { 'tsquery.lua' },
     files = {
       'runtime/lua/vim/treesitter/_meta/',
       'runtime/lua/vim/treesitter.lua',

src/nvim/lua/treesitter.c

@@ -1491,6 +1491,8 @@ static struct luaL_Reg query_meta[] = {
   { "__gc", query_gc },
   { "__tostring", query_tostring },
   { "inspect", query_inspect },
+  { "disable_capture", query_disable_capture },
+  { "disable_pattern", query_disable_pattern },
   { NULL, NULL }
 };
 
@@ -1689,6 +1691,23 @@ static int query_inspect(lua_State *L)
   return 1;
 }
 
+static int query_disable_capture(lua_State *L)
+{
+  TSQuery *query = query_check(L, 1);
+  size_t name_len;
+  const char *name = luaL_checklstring(L, 2, &name_len);
+  ts_query_disable_capture(query, name, (uint32_t)name_len);
+  return 0;
+}
+
+static int query_disable_pattern(lua_State *L)
+{
+  TSQuery *query = query_check(L, 1);
+  const uint32_t pattern_index = (uint32_t)luaL_checkinteger(L, 2);
+  ts_query_disable_pattern(query, pattern_index - 1);
+  return 0;
+}
+
 // Library init
 
 static void build_meta(lua_State *L, const char *tname, const luaL_Reg *meta)

test/functional/treesitter/highlight_spec.lua

@@ -640,8 +640,14 @@ describe('treesitter highlighting (C)', function()
     }
 
     eq({
-      { capture = 'constant', metadata = { priority = '101' }, lang = 'c', id = 14 },
-      { capture = 'type', metadata = {}, lang = 'c', id = 3 },
+      {
+        capture = 'constant',
+        metadata = { priority = '101' },
+        lang = 'c',
+        id = 14,
+        pattern_id = 23,
+      },
+      { capture = 'type', metadata = {}, lang = 'c', id = 3, pattern_id = 16 },
     }, exec_lua [[ return vim.treesitter.get_captures_at_pos(0, 0, 2) ]])
   end)
 

test/functional/treesitter/query_spec.lua

@@ -912,4 +912,52 @@ void ui_refresh(void)
       eq({ 2, { 1, 1, 2, 2 } }, result)
     end)
   end)
+
+  describe('TSQuery', function()
+    local source = [[
+      void foo(int x, int y);
+    ]]
+
+    local query_text = [[
+      ((identifier) @func
+        (#eq? @func "foo"))
+      ((identifier) @param
+        (#eq? @param "x"))
+      ((identifier) @param
+        (#eq? @param "y"))
+    ]]
+
+    ---@param query string
+    ---@param disabled { capture: string?, pattern: integer? }
+    local function get_patterns(query, disabled)
+      local q = vim.treesitter.query.parse('c', query)
+      if disabled.capture then
+        q.query:disable_capture(disabled.capture)
+      end
+      if disabled.pattern then
+        q.query:disable_pattern(disabled.pattern)
+      end
+
+      local parser = vim.treesitter.get_parser(0, 'c')
+      local root = parser:parse()[1]:root()
+      local captures = {} ---@type {id: number, pattern: number}[]
+      for id, _, _, match in q:iter_captures(root, 0) do
+        local _, pattern = match:info()
+        captures[#captures + 1] = { id = id, pattern = pattern }
+      end
+      return captures
+    end
+
+    it('supports disabling patterns', function()
+      insert(source)
+      local result = exec_lua(get_patterns, query_text, { pattern = 2 })
+      eq({ { id = 1, pattern = 1 }, { id = 2, pattern = 3 } }, result)
+    end)
+
+    it('supports disabling captures', function()
+      insert(source)
+      local result = exec_lua(get_patterns, query_text, { capture = 'param' })
+      eq({ { id = 1, pattern = 1 } }, result)
+    end)
+  end)
 end)