Feat: use treesitter to determine filetype at cursor position

README.md

@@ -8,6 +8,11 @@
   <img alt="Screenshot" title="cmp-nvim-ultisnips" src="screenshots/preview.png" width="80%" height="80%">
 </p>
 
+## Features
+- **Composable Mappings**: get rid of boilerplate code in your config
+- **Customization**: change how your snippets are displayed by cmp
+- **Treesitter Integration**: show snippets based on your current position in a file
+
 ## Installation and Recommended Mappings
 
 Check out the [Mappings](#Mappings) section if you want to define your own mappings.
@@ -157,4 +162,4 @@ use({
     vim.g.UltiSnipsRemoveSelectModeMappings = 0
   end,
 })
-```
+```

autoload/cmp_nvim_ultisnips.vim

@@ -1,3 +1,5 @@
+" TODO: move python code into separate files
+
 " Retrieves additional snippet information that is not directly accessible
 " using the UltiSnips API functions. Returns a list of tables (one table
 " per snippet) with the keys "trigger", "description", "options" and "value".
@@ -9,8 +11,7 @@ pythonx << EOF
 import vim
 from UltiSnips import UltiSnips_Manager, vim_helper
 
-expandable_only = vim.eval("a:expandable_only")
-if expandable_only == "True":
+if vim.eval("a:expandable_only") == "True":
     before = vim_helper.buf.line_till_cursor
     snippets = UltiSnips_Manager._snips(before, True)
 else:
@@ -31,6 +32,42 @@ EOF
 return g:_cmpu_current_snippets
 endfunction
 
+function cmp_nvim_ultisnips#update_filetypes(remove_filetype, additional_filetype)
+pythonx << EOF
+import vim
+from UltiSnips import UltiSnips_Manager, vim_helper
+
+remove_ft = vim.eval("a:remove_filetype")
+additional_ft = vim.eval("a:additional_filetype")
+cur_fts = UltiSnips_Manager._added_buffer_filetypes[vim_helper.buf.number]
+
+if remove_ft in cur_fts:
+  cur_fts.remove(remove_ft)
+if additional_ft not in cur_fts:
+  cur_fts.append(additional_ft)
+EOF
+endfunction
+
+function! cmp_nvim_ultisnips#remove_snippet_filetype(filetype)
+pythonx << EOF
+import vim
+from UltiSnips import UltiSnips_Manager, vim_helper
+
+ft = vim.eval("a:filetype")
+cur_fts = UltiSnips_Manager._added_buffer_filetypes[vim_helper.buf.number]
+if ft in cur_fts:
+  cur_fts.remove(ft)
+EOF
+endfunction
+
+function! cmp_nvim_ultisnips#setup_treesitter_autocmds()
+  augroup cmp_nvim_ultisnips
+    autocmd!
+    autocmd CursorMovedI * lua require("cmp_nvim_ultisnips.treesitter").update_ft_at_cursor()
+    autocmd InsertLeave * lua require("cmp_nvim_ultisnips.treesitter").clear_ft_at_cursor()
+  augroup end
+endfunction
+
 " Define silent mappings
 imap <silent> <Plug>(cmpu-expand)
 \ <C-r>=[UltiSnips#CursorMoved(), UltiSnips#ExpandSnippet()][1]<cr>

lua/cmp_nvim_ultisnips/config.lua

@@ -5,6 +5,7 @@ local M = {}
 M.default_config = {
   show_snippets = "expandable",
   documentation = cmpu_snippets.documentation,
+  filetype_source = "filename", -- TODO: better name for filename?
 }
 
 function M.get_user_config(config)
@@ -18,6 +19,15 @@ function M.get_user_config(config)
       "either 'expandable' or 'all'",
     },
     documentation = { user_config.documentation, "function" },
+    filetype_source = {
+      user_config.filetype_source,
+      function(arg)
+        return arg == "filename" or arg == "treesitter" or arg == "both"
+      end,
+      -- currently, only 'filename' and 'both' work;
+      -- TODO: find way to support filetype_source=treesitter only
+      "either 'filename' or 'treesitter' or 'both'",
+    },
   })
   return user_config
 end

lua/cmp_nvim_ultisnips/source.lua

@@ -6,6 +6,9 @@ function source.new(config)
   local self = setmetatable({}, { __index = source })
   self.config = config
   self.expandable_only = config.show_snippets == "expandable"
+  if config.filetype_source == "treesitter" or config.filetype_source == "both" then
+    vim.fn["cmp_nvim_ultisnips#setup_treesitter_autocmds"]()
+  end
   return self
 end
 

lua/cmp_nvim_ultisnips/treesitter.lua

@@ -0,0 +1,43 @@
+local M = {}
+
+-- Check if treesitter is available
+local ok_parsers, ts_parsers = pcall(require, "nvim-treesitter.parsers")
+if not ok_parsers then
+  ts_parsers = nil
+end
+
+local ok_utils, ts_utils = pcall(require, "nvim-treesitter.ts_utils")
+if not ok_utils then
+  ts_utils = nil
+end
+
+local function get_ft_at_cursor()
+  if not ok_parsers or not ok_utils then
+    error("[cmp_nvim_ultisnips.treesitter] Please install the nvim-treesitter plugin to use this feature.")
+  end
+
+  local parser = ts_parsers.get_parser()
+  local cur_node = ts_utils.get_node_at_cursor()
+  if cur_node then
+    return parser:language_for_range({ cur_node:range() }):lang()
+  end
+  return nil
+end
+
+local cur_ft_at_cursor
+function M.update_ft_at_cursor()
+  local new_lang = get_ft_at_cursor()
+  if new_lang ~= cur_ft_at_cursor then
+    vim.fn["cmp_nvim_ultisnips#update_filetypes"](cur_ft_at_cursor, new_lang)
+    cur_ft_at_cursor = new_lang
+  end
+end
+
+function M.clear_ft_at_cursor()
+  if cur_ft_at_cursor ~= nil and cur_ft_at_cursor ~= vim.bo.filetype then
+    vim.fn["cmp_nvim_ultisnips#remove_snippet_filetype"](cur_ft_at_cursor)
+    cur_ft_at_cursor = nil
+  end
+end
+
+return M