feat(treesitter): add lang parameter to the query editor (#25181)

4 files changed, 20 insertions, 13 deletions

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index ee34c45cce..e1e280c4db 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -853,7 +853,7 @@ add_predicate({name}, {handler}, {force})
                      meanings
       • {force}    (boolean|nil)
 
-edit()                                           *vim.treesitter.query.edit()*
+edit({lang})                                     *vim.treesitter.query.edit()*
     Open a window for live editing of a treesitter query.
 
     Can also be shown with `:EditQuery`.                          *:EditQuery*
@@ -861,6 +861,10 @@ edit()                                           *vim.treesitter.query.edit()*
     Note that the editor opens a scratch buffer, and so queries aren't
     persisted on disk.
 
+    Parameters: ~
+      • {lang}  (string|nil) language to open the query editor for. If
+                omitted, inferred from the current buffer's filetype.
+
 get({lang}, {query_name})                         *vim.treesitter.query.get()*
     Returns the runtime query {query_name} for {lang}.
 
diff --git a/runtime/lua/vim/treesitter/dev.lua b/runtime/lua/vim/treesitter/dev.lua
index bc54853103..7f24ba8590 100644
--- a/runtime/lua/vim/treesitter/dev.lua
+++ b/runtime/lua/vim/treesitter/dev.lua
@@ -468,11 +468,11 @@ local edit_ns = api.nvim_create_namespace('treesitter/dev-edit')
 
 ---@param query_win integer
 ---@param base_win integer
-local function update_editor_highlights(query_win, base_win)
+---@param lang string
+local function update_editor_highlights(query_win, base_win, lang)
   local base_buf = api.nvim_win_get_buf(base_win)
   local query_buf = api.nvim_win_get_buf(query_win)
-  local parser = vim.treesitter.get_parser(base_buf)
-  local lang = parser:lang()
+  local parser = vim.treesitter.get_parser(base_buf, lang)
   api.nvim_buf_clear_namespace(base_buf, edit_ns, 0, -1)
   local query_content = table.concat(api.nvim_buf_get_lines(query_buf, 0, -1, false), '\n')
 
@@ -506,7 +506,8 @@ local function update_editor_highlights(query_win, base_win)
 end
 
 --- @private
-function M.edit_query()
+--- @param lang? string language to open the query editor for.
+function M.edit_query(lang)
   local buf = api.nvim_get_current_buf()
   local win = api.nvim_get_current_win()
 
@@ -528,11 +529,11 @@ function M.edit_query()
   end
   vim.cmd(cmd)
 
-  local ok, parser = pcall(vim.treesitter.get_parser, buf)
+  local ok, parser = pcall(vim.treesitter.get_parser, buf, lang)
   if not ok then
     return nil, 'No parser available for the given buffer'
   end
-  local lang = parser:lang()
+  lang = parser:lang()
 
   local query_win = api.nvim_get_current_win()
   local query_buf = api.nvim_win_get_buf(query_win)
@@ -561,7 +562,7 @@ function M.edit_query()
     desc = 'Update query editor highlights when the cursor moves',
     callback = function()
       if api.nvim_win_is_valid(win) then
-        update_editor_highlights(query_win, win)
+        update_editor_highlights(query_win, win, lang)
       end
     end,
   })
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index e83ad00eeb..44ed37d64e 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -848,8 +848,10 @@ end
 --- Can also be shown with `:EditQuery`. *:EditQuery*
 ---
 --- Note that the editor opens a scratch buffer, and so queries aren't persisted on disk.
-function M.edit()
-  require('vim.treesitter.dev').edit_query()
+---
+--- @param lang? string language to open the query editor for. If omitted, inferred from the current buffer's filetype.
+function M.edit(lang)
+  require('vim.treesitter.dev').edit_query(lang)
 end
 
 return M
diff --git a/runtime/plugin/nvim.lua b/runtime/plugin/nvim.lua
index 13a54a5e20..743d3044b6 100644
--- a/runtime/plugin/nvim.lua
+++ b/runtime/plugin/nvim.lua
@@ -19,6 +19,6 @@ vim.api.nvim_create_user_command('InspectTree', function(cmd)
   end
 end, { desc = 'Inspect treesitter language tree for buffer', count = true })
 
-vim.api.nvim_create_user_command('EditQuery', function()
-  vim.treesitter.query.edit()
-end, { desc = 'Edit treesitter query' })
+vim.api.nvim_create_user_command('EditQuery', function(cmd)
+  vim.treesitter.query.edit(cmd.fargs[1])
+end, { desc = 'Edit treesitter query', nargs = '?' })