feat: Add TSEditQuery

3 files changed, 44 insertions, 0 deletions

diff --git a/autoload/nvim_treesitter.vim b/autoload/nvim_treesitter.vim
index 2b02db2d8..1216e9967 100644
--- a/autoload/nvim_treesitter.vim
+++ b/autoload/nvim_treesitter.vim
@@ -18,6 +18,10 @@ function! nvim_treesitter#available_modules(arglead, cmdline, cursorpos) abort
   return join(luaeval("require'nvim-treesitter.configs'.available_modules()"), "\n")
 endfunction
 
+function! nvim_treesitter#available_query_groups(arglead, cmdline, cursorpos) abort
+  return join(luaeval("require'nvim-treesitter.query'.available_query_groups()"), "\n")
+endfunction
+
 function! nvim_treesitter#indent() abort
 	return luaeval(printf('require"nvim-treesitter.indent".get_indent(%d)', v:lnum))
 endfunction
diff --git a/lua/nvim-treesitter/configs.lua b/lua/nvim-treesitter/configs.lua
index 973b9cb6d..76a495b9a 100644
--- a/lua/nvim-treesitter/configs.lua
+++ b/lua/nvim-treesitter/configs.lua
@@ -1,6 +1,7 @@
 local api = vim.api
 
 local queries = require'nvim-treesitter.query'
+local nvim_query = require'vim.treesitter.query'
 local parsers = require'nvim-treesitter.parsers'
 local utils = require'nvim-treesitter.utils'
 local caching = require'nvim-treesitter.caching'
@@ -196,6 +197,25 @@ local function config_info(process_function)
   print(vim.inspect(config, {process = process_function}))
 end
 
+function M.edit_query_file(query_group, lang)
+  lang = lang or parsers.get_buf_lang()
+  local files = nvim_query.get_query_files(lang, query_group, true)
+  if #files == 0 then
+    local folder = utils.join_path(vim.fn.stdpath('config'), 'after', 'queries', lang)
+    local file = utils.join_path(folder, query_group..'.scm')
+    pcall(vim.fn.mkdir, folder, "p", "0755")
+    vim.cmd(':edit '..file)
+  elseif #files == 1 then
+    vim.cmd(':edit '..files[1])
+  else
+    local counter = 0
+    local choice = vim.fn.inputlist(vim.tbl_map(function(f) counter = counter + 1;return counter..'. '..f end, files))
+    if choice > 0 then
+      vim.cmd(':edit '..files[choice])
+    end
+  end
+end
+
 M.commands = {
   TSBufEnable = {
     run = enable_module,
@@ -245,6 +265,13 @@ M.commands = {
       "-nargs=0",
     },
   },
+  TSEditQuery = {
+    run = M.edit_query_file,
+    args = {
+      "-nargs=+",
+      "-complete=custom,nvim_treesitter#available_query_groups",
+    },
+  },
 }
 
 -- @param mod: module (string)
diff --git a/lua/nvim-treesitter/query.lua b/lua/nvim-treesitter/query.lua
index 9c929defa..07896f22f 100644
--- a/lua/nvim-treesitter/query.lua
+++ b/lua/nvim-treesitter/query.lua
@@ -23,6 +23,19 @@ for _, query in ipairs(M.built_in_query_groups) do
   M["has_" .. query] = get_query_guard(query)
 end
 
+function M.available_query_groups()
+  local query_files = api.nvim_get_runtime_file('queries/*/*.scm', true)
+  local groups = {}
+  for _, f in ipairs(query_files) do
+    groups[vim.fn.fnamemodify(f, ':t:r')] = true
+  end
+  local list = {}
+  for k, _ in pairs(groups) do
+    table.insert(list, k)
+  end
+  return list
+end
+
 do
   local query_cache = caching.create_buffer_cache()