Merge pull request #14252 from tjdevries/tjdevries/ts_query_overrides

ts: Add per-language query overriding

runtime/lua/vim/treesitter/query.lua

@@ -90,6 +90,27 @@ local function new_match_metadata()
   return setmetatable({}, match_metatable)
 end
 
+--- The explicitly set queries from |vim.treesitter.query.set_query()|
+local explicit_queries = setmetatable({}, {
+  __index = function(t, k)
+    local lang_queries = {}
+    rawset(t, k, lang_queries)
+
+    return lang_queries
+  end,
+})
+
+--- Sets the runtime query {query_name} for {lang}
+---
+--- This allows users to override any runtime files and/or configuration
+--- set by plugins.
+---@param lang string: The language to use for the query
+---@param query_name string: The name of the query (i.e. "highlights")
+---@param text string: The query text (unparsed).
+function M.set_query(lang, query_name, text)
+  explicit_queries[lang][query_name] = M.parse_query(lang, text)
+end
+
 --- Returns the runtime query {query_name} for {lang}.
 --
 -- @param lang The language to use for the query
@@ -97,6 +118,10 @@ end
 --
 -- @return The corresponding query, parsed.
 function M.get_query(lang, query_name)
+  if explicit_queries[lang][query_name] then
+    return explicit_queries[lang][query_name]
+  end
+
   local query_files = M.get_query_files(lang, query_name)
   local query_string = read_query_files(query_files)
 

test/functional/treesitter/highlight_spec.lua

@@ -472,4 +472,45 @@ describe('treesitter highlighting', function()
                                                                        |
     ]]}
   end)
+
+  it("supports overriding queries, like ", function()
+    if pending_c_parser(pending) then return end
+
+    insert([[
+    int x = INT_MAX;
+    #define READ_STRING(x, y) (char_u *)read_string((x), (size_t)(y))
+    #define foo void main() { \
+                  return 42;  \
+                }
+    ]])
+
+    exec_lua [[
+      local injection_query = "(preproc_def (preproc_arg) @c) (preproc_function_def value: (preproc_arg) @c)"
+      require('vim.treesitter.query').set_query("c", "highlights", hl_query)
+      require('vim.treesitter.query').set_query("c", "injections", injection_query)
+
+      vim.treesitter.highlighter.new(vim.treesitter.get_parser(0, "c"))
+    ]]
+
+    screen:expect{grid=[[
+      {3:int} x = {5:INT_MAX};                                                 |
+      #define {5:READ_STRING}(x, y) ({3:char_u} *)read_string((x), ({3:size_t})(y))|
+      #define foo {3:void} main() { \                                      |
+                    {4:return} {5:42};  \                                      |
+                  }                                                    |
+      ^                                                                 |
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+                                                                       |
+    ]]}
+  end)
 end)