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

4 files changed, 33 insertions, 0 deletions

diff --git a/README.md b/README.md
index fc0098706..04d72554d 100644
--- a/README.md
+++ b/README.md
@@ -98,6 +98,9 @@ require'nvim-treesitter.configs'.setup {
     highlight = {
         enable = true,                    -- false will disable the whole extension
         disable = { 'c', 'rust' },        -- list of language that will be disabled
+        custom_captures = {               -- mapping of user defined captures to highlight groups
+          -- ["foo.bar"] = "Identifier"   -- highlight own capture @foo.bar with highlight group "Identifier", see :h nvim-treesitter-query-extensions
+        },            
     },
     incremental_selection = {
         enable = true,
diff --git a/doc/nvim-treesitter.txt b/doc/nvim-treesitter.txt
index c9adf93b2..5b619f8c4 100644
--- a/doc/nvim-treesitter.txt
+++ b/doc/nvim-treesitter.txt
@@ -33,6 +33,9 @@ By default, everything is disabled. To enable support for features, in your `ini
     highlight = {
         enable = true,                    -- false will disable the whole extension
         disable = { 'c', 'rust' },        -- list of language that will be disabled
+        custom_captures = {               -- mapping of user defined captures to highlight groups
+          -- ["foo.bar"] = "Identifier"   -- highlight own capture @foo.bar with highlight group "Identifier", see :h nvim-treesitter-query-extensions
+        },            
     },
     incremental_selection = {
         enable = true,
@@ -48,6 +51,18 @@ By default, everything is disabled. To enable support for features, in your `ini
   }
   EOF
 <
+==============================================================================
+USER QUERY EXTENSIONS                        *nvim-treesitter-query-extensions*
+
+You can add your own query files by placing a query file in vim's runtime path
+after `nvim-treesitter` is sourced.  If the language has a built in query file,
+that file will be appended to or it will be used (useful for languages not yet
+supported).
+
+For example, you can add files to `<vim-config-dir>/after/queries/lua/highlights.scm`
+to add more queries to lua highlights.  You can also manually add query paths
+to the runtime path by adding this to your vim config `set rtp+='path/to/queries'`.
+
 
 ==============================================================================
 COMMANDS                                          *nvim-treesitter-commands*
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