feat: add locals to setup procedure

2 files changed, 19 insertions, 5 deletions

diff --git a/lua/nvim-treesitter.lua b/lua/nvim-treesitter.lua
index abb3b6d91..ef6c31506 100644
--- a/lua/nvim-treesitter.lua
+++ b/lua/nvim-treesitter.lua
@@ -1,6 +1,7 @@
 local api = vim.api
 local parsers = require'nvim-treesitter.parsers'
 local install = require'nvim-treesitter.install'
+local locals = require'nvim-treesitter.locals'
 
 local M = {}
 
@@ -8,6 +9,12 @@ local M = {}
 -- this is the main interface through the plugin
 function M.setup(lang)
   if parsers.has_parser(lang) then
+    if locals.is_supported(lang) then
+      print("Locals setup for", lang)
+      api.nvim_command(string.format([[
+      autocmd NvimTreesitter FileType %s lua vim.api.nvim_buf_attach(0, true, {on_lines=require'nvim-treesitter.locals'.on_lines})
+      ]], lang))
+    end
   end
 end
 
diff --git a/lua/nvim-treesitter/locals.lua b/lua/nvim-treesitter/locals.lua
index cd8a3e306..cd0c268db 100644
--- a/lua/nvim-treesitter/locals.lua
+++ b/lua/nvim-treesitter/locals.lua
@@ -10,9 +10,12 @@ local M = {
   locals={}
 }
 
+function M.is_supported(lang)
+  return queries.get_query(lang, "locals") ~= nil
+end
+
 function M.collect_locals(bufnr)
   local ft = api.nvim_buf_get_option(bufnr, "ft")
-  local query = queries.get_query(ft, "locals")
 
   if ft then
     local query = queries.get_query(ft, 'locals')
@@ -33,12 +36,16 @@ function M.collect_locals(bufnr)
   end
 end
 
-function M.on_lines(_, buf, _, firstline, lastline, new_lastline, _)
+function M.on_lines(_, buf, _, firstline, lastline, new_lastline)
   M.locals[buf] = M.collect_locals(buf)
 end
 
+function M.get_locals(bufnr)
+  return M.locals[bufnr or api.nvim_get_current_buf()] or {}
+end
+
 function M.get_definitions(bufnr)
-  local locals = M.locals[bufnr]
+  local locals = M.get_locals(bufnr)
 
   local defs = {}
 
@@ -52,7 +59,7 @@ function M.get_definitions(bufnr)
 end
 
 function M.get_scopes(bufnr)
-  local locals = M.locals[bufnr]
+  local locals = M.get_locals(bufnr)
 
   local scopes = {}
 
@@ -66,7 +73,7 @@ function M.get_scopes(bufnr)
 end
 
 function M.get_references(bufnr)
-  local locals = M.locals[bufnr or api.nvim_get_current_buf()]
+  local locals = M.get_locals(bufnr)
 
   local refs = {}