feat: add omnifunc for query filetype

after/ftplugin/query.vim

@@ -0,0 +1 @@
+setlocal omnifunc=v:lua.require'nvim_treesitter_playground_query_omnifunc'.omnifunc

lua/nvim-treesitter-playground/query_linter.lua

@@ -50,18 +50,23 @@ local function query_lang_from_playground_buf(buf)
   end
 end
 
-function M.lint(query_buf)
-  query_buf = query_buf or api.nvim_get_current_buf()
-  M.clear_virtual_text(query_buf)
-  M.lints[query_buf] = {}
-
-  local filename = api.nvim_buf_get_name(query_buf)
+function M.guess_query_lang(buf)
+  local filename = api.nvim_buf_get_name(buf)
   local ok, query_lang = pcall(vim.fn.fnamemodify, filename, ":p:h:t")
   query_lang = filename ~= "" and query_lang
   query_lang = ok and query_lang
   if not query_lang then
-    query_lang = query_lang_from_playground_buf(query_buf)
+    query_lang = query_lang_from_playground_buf(buf)
   end
+  return query_lang
+end
+
+function M.lint(query_buf)
+  query_buf = query_buf or api.nvim_get_current_buf()
+  M.clear_virtual_text(query_buf)
+  M.lints[query_buf] = {}
+
+  local query_lang = M.guess_query_lang(query_buf)
 
   local ok, parser_info = pcall(vim.treesitter.inspect_language, query_lang)
 

lua/nvim_treesitter_playground_query_omnifunc.lua

@@ -0,0 +1,49 @@
+local query_linter = require 'nvim-treesitter-playground.query_linter'
+local tsc = require 'vim.treesitter.query'
+
+local M = {}
+
+function M.omnifunc(findstart, base)
+  if findstart == 1 then
+    local start = vim.fn.col('.') - 1
+    local result = vim.fn.matchstrpos(vim.fn.getline('.'):sub(1, start), '\\v(["#\\-]|\\w)*$')
+    return result[2]
+  end
+
+  local buf = vim.api.nvim_get_current_buf()
+  local query_lang = query_linter.guess_query_lang(buf)
+
+  local ok, parser_info = pcall(vim.treesitter.inspect_language, query_lang)
+
+  if ok then
+    local items = {}
+    for _, f in pairs(parser_info.fields) do
+      if f:find(base, 1, true) == 1 then
+        table.insert(items, f..':')
+      end
+    end
+    for _, p in pairs(tsc.list_predicates()) do
+      local text = '#'..p
+      local found = text:find(base, 1, true)
+      if found and found <= 2 then -- with or without '#'
+        table.insert(items, text)
+      end
+      text = '#not-'..p
+      found = text:find(base, 1, true)
+      if found and found <= 2 then -- with or without '#'
+        table.insert(items, text)
+      end
+    end
+    for _, s in pairs(parser_info.symbols) do
+      local text = s[2] and s[1] or '"'..vim.fn.escape(s[1], '\\')..'"'
+      if text:find(base, 1, true) == 1 then
+        table.insert(items, text)
+      end
+    end
+    return {words = items, refresh= 'always'}
+  else
+    return -2
+  end
+end
+
+return M