Fix #167: Add custom_captures config key to set highlights for custom queries

2 files changed, 15 insertions, 0 deletions

diff --git a/lua/nvim-treesitter/configs.lua b/lua/nvim-treesitter/configs.lua
index 85ce0f632..abf739dd9 100644
--- a/lua/nvim-treesitter/configs.lua
+++ b/lua/nvim-treesitter/configs.lua
@@ -18,6 +18,7 @@ local builtin_modules = {
     module_path = 'nvim-treesitter.highlight',
     enable = false,
     disable = {},
+    custom_captures = {},
     is_supported = queries.has_highlights
   },
   incremental_selection = {
@@ -61,6 +62,8 @@ local builtin_modules = {
   }
 }
 
+local special_config_keys = {'enable', 'disable', 'keymaps'}
+
 local M = {}
 
 -- Resolves a module by requiring the `module_path` or using the module definition.
@@ -276,6 +279,12 @@ function M.setup_module(mod, data, mod_name)
         end
       end
     end
+    for k, v in pairs(data) do
+      -- Just copy all non-special configuration keys
+      if not vim.tbl_contains(special_config_keys, k) then
+        mod[k] = v
+      end
+    end
   elseif type(data) == 'table' and type(mod) == 'table' then
     for key, value in pairs(mod) do
       M.setup_module(value, data[key], key)
diff --git a/lua/nvim-treesitter/highlight.lua b/lua/nvim-treesitter/highlight.lua
index 1b4722435..0d4f3f547 100644
--- a/lua/nvim-treesitter/highlight.lua
+++ b/lua/nvim-treesitter/highlight.lua
@@ -3,6 +3,7 @@ local ts = vim.treesitter
 
 local queries = require'nvim-treesitter.query'
 local parsers = require'nvim-treesitter.parsers'
+local configs = require'nvim-treesitter.configs'
 
 local M = {
   highlighters = {}
@@ -54,6 +55,11 @@ hlmap["include"] = "TSInclude"
 function M.attach(bufnr, lang)
   local bufnr = bufnr or api.nvim_get_current_buf()
   local lang = lang or parsers.ft_to_lang(api.nvim_buf_get_option(bufnr, 'ft'))
+  local config = configs.get_module('highlight')
+
+  for k, v in pairs(config.custom_captures) do
+    hlmap[k] = v
+  end
 
   local query = queries.get_query(lang, "highlights")
   if not query then return end