From e6b743ccd4e780bc9cd85b707de67df72eea1a23 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 28 Aug 2023 09:50:34 +0100
Subject: [PATCH] refactor: autocmds and setup

---
 lua/treesitter-context.lua    | 62 +++++++++++++----------------------
 plugin/treesitter-context.lua |  1 +
 plugin/treesitter-context.vim |  1 -
 3 files changed, 24 insertions(+), 40 deletions(-)
 create mode 100644 plugin/treesitter-context.lua
 delete mode 100644 plugin/treesitter-context.vim

diff --git a/lua/treesitter-context.lua b/lua/treesitter-context.lua
index b1cca17..8745987 100644
--- a/lua/treesitter-context.lua
+++ b/lua/treesitter-context.lua
@@ -765,34 +765,23 @@ local update = throttle(function()
   end
 end)
 
---- @param group string
---- @return function
-local function autocmd_for_group(group)
-  local gid = augroup(group, {})
-  return function(event, opts)
-    ---@diagnostic disable:no-unknown
-    if opts then
-      if type(opts) == 'function' then
-        opts = { callback = opts }
-      elseif opts[1] then
-        opts.callback = opts[1]
-        opts[1] = nil
-      end
-    else
-      opts = {}
-    end
-    opts.group = gid
-    api.nvim_create_autocmd(event, opts)
-  end
-end
-
 local M = {
   config = config,
 }
 
-function M.enable()
-  local autocmd = autocmd_for_group('treesitter_context_update')
+local group = augroup('treesitter_context_update', {})
 
+---@param event string|string[]
+---@param callback fun(args: table)
+---@param opts? vim.api.keyset.create_autocmd
+local function autocmd(event, callback, opts)
+  opts = opts or {}
+  opts.callback = callback
+  opts.group = group
+  api.nvim_create_autocmd(event, opts)
+end
+
+function M.enable()
   local cbuf = api.nvim_get_current_buf()
 
   attached[cbuf] = true
@@ -824,8 +813,8 @@ function M.enable()
 
   autocmd({ 'BufLeave', 'WinLeave' }, close)
 
-  autocmd('User', { close, pattern = 'SessionSavePre' })
-  autocmd('User', { update, pattern = 'SessionSavePost' })
+  autocmd('User', close, { pattern = 'SessionSavePre' })
+  autocmd('User', update, { pattern = 'SessionSavePost' })
 
   update()
   enabled = true
@@ -863,6 +852,15 @@ function M.setup(options)
   else
     M.disable()
   end
+
+  command('TSContextEnable', M.enable, {})
+  command('TSContextDisable', M.disable, {})
+  command('TSContextToggle', M.toggle, {})
+
+  api.nvim_set_hl(0, 'TreesitterContext', { link = 'NormalFloat', default = true })
+  api.nvim_set_hl(0, 'TreesitterContextLineNumber', { link = 'LineNr', default = true })
+  api.nvim_set_hl(0, 'TreesitterContextBottom', { link = 'NONE', default = true })
+  api.nvim_set_hl(0, 'TreesitterContextSeparator', { link = 'FloatBorder', default = true })
 end
 
 function M.go_to_context()
@@ -885,18 +883,4 @@ function M.go_to_context()
   api.nvim_win_set_cursor(0, { context[1] + 1, context[2] })
 end
 
-command('TSContextEnable', M.enable, {})
-command('TSContextDisable', M.disable, {})
-command('TSContextToggle', M.toggle, {})
-
-api.nvim_set_hl(0, 'TreesitterContext', { link = 'NormalFloat', default = true })
-api.nvim_set_hl(0, 'TreesitterContextLineNumber', { link = 'LineNr', default = true })
-api.nvim_set_hl(0, 'TreesitterContextBottom', { link = 'NONE', default = true })
-api.nvim_set_hl(0, 'TreesitterContextSeparator', { link = 'FloatBorder', default = true })
-
--- Setup with default options if user didn't call setup()
-autocmd_for_group('treesitter_context')('VimEnter', function()
-  M.setup()
-end)
-
 return M
diff --git a/plugin/treesitter-context.lua b/plugin/treesitter-context.lua
new file mode 100644
index 0000000..da358bb
--- /dev/null
+++ b/plugin/treesitter-context.lua
@@ -0,0 +1 @@
+require('treesitter-context').setup()
diff --git a/plugin/treesitter-context.vim b/plugin/treesitter-context.vim
deleted file mode 100644
index 133f6ff..0000000
--- a/plugin/treesitter-context.vim
+++ /dev/null
@@ -1 +0,0 @@
-lua require'treesitter-context'