Feat: integrate 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
+- **Treesitter Integration**: show snippets based on the filetype at your cursor position
+- **Customization**: change which and how snippets are displayed by cmp
+
 ## Installation and Recommended Mappings
 
 Check out the [Mappings](#Mappings) section if you want to define your own mappings.
@@ -20,7 +25,9 @@ use({
     config = function()
       -- optional call to setup (see customization section)
       require("cmp_nvim_ultisnips").setup{}
-    end
+    end,
+    -- If you want to enable filetype detection based on treesitter:
+    -- requires = { "nvim-treesitter/nvim-treesitter" },
   },
   config = function()
     local cmp_ultisnips_mappings = require("cmp_nvim_ultisnips.mappings")
@@ -96,6 +103,18 @@ Then the `command` function is run. If none match, `fallback` is called.
 
 ### Available Options
 
+`filetype_source: "treesitter" | "ultisnips_default"`
+
+Determines how the filetype of a buffer is identified. This option affects which snippets are available by UltiSnips.
+If set to `"treesitter"` and the [`nvim-treesitter`](https://github.com/nvim-treesitter/nvim-treesitter) plugin is installed, only snippets
+that match the filetype at the current cursor position are shown (as well as snippets included via UltiSnips' `extends` directive). Otherwise, or if
+treesitter could not determine the filetype at the current position, the available snippets
+are handled entirely by UltiSnips.
+
+**Default:** `"treesitter"`
+
+---
+
 `show_snippets: "expandable" | "all"`
 
 If set to `"expandable"`, only those snippets currently expandable by UltiSnips will be
@@ -135,6 +154,7 @@ Note: calling the setup function is only required if you wish to customize this
 
 ```lua
 require("cmp_nvim_ultisnips").setup {
+  filetype_source = "treesitter",
   show_snippets = "all",
   documentation = function(snippet)
     return snippet.description
@@ -143,12 +163,13 @@ require("cmp_nvim_ultisnips").setup {
 ```
 
 ## Credit
-[Compe source for UltiSnips](https://github.com/hrsh7th/nvim-compe/blob/master/lua/compe_ultisnips/init.lua)
+- [Compe source for UltiSnips](https://github.com/hrsh7th/nvim-compe/blob/master/lua/compe_ultisnips/init.lua)
+- The Treesitter integration was inspired by [this Luasnip PR](https://github.com/L3MON4D3/LuaSnip/pull/226)
 
 ## Known Issues
 
 UltiSnips was auto-removing tab mappings for select mode, that way it was not possible to jump through snippet stops.
-We have to disable this by setting `UltiSnipsRemoveSelectModeMappings = 0` (Credit [JoseConseco](https://github.com/quangnguyen30192/cmp-nvim-ultisnips/issues/5))
+We have to disable this by setting `UltiSnipsRemoveSelectModeMappings = 0` (credit to [JoseConseco](https://github.com/quangnguyen30192/cmp-nvim-ultisnips/issues/5)).
 ```lua
 use({
   "SirVer/ultisnips",
@@ -157,4 +178,4 @@ use({
     vim.g.UltiSnipsRemoveSelectModeMappings = 0
   end,
 })
-```
+```

autoload/cmp_nvim_ultisnips.vim

@@ -1,36 +1,69 @@
+" 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".
 "
 " If 'expandable_only' is 1, only expandable snippets are returned, otherwise all
 " snippets for the current filetype are returned.
 function! cmp_nvim_ultisnips#get_current_snippets(expandable_only)
-pythonx << EOF
+python3 << 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:
     snippets = UltiSnips_Manager._snips("", True)
 
-snippets_info = []
 vim.command('let g:_cmpu_current_snippets = []')
 for snippet in snippets:
     vim.command(
-      "call add(g:_cmpu_current_snippets, {"\
-        "'trigger': pyxeval('str(snippet._trigger)'),"\
-        "'description': pyxeval('str(snippet._description)'),"\
-        "'options': pyxeval('str(snippet._opts)'),"\
-        "'value': pyxeval('str(snippet._value)'),"\
+      "call add(g:_cmpu_current_snippets, {"
+      "'trigger': py3eval('str(snippet._trigger)'),"
+      "'description': py3eval('str(snippet._description)'),"
+      "'options': py3eval('str(snippet._opts)'),"
+      "'value': py3eval('str(snippet._value)'),"
       "})"
     )
 EOF
 return g:_cmpu_current_snippets
 endfunction
 
+function cmp_nvim_ultisnips#set_filetype(filetype)
+python3 << EOF
+import vim
+from UltiSnips import vim_helper
+
+filetype = vim.eval("a:filetype")
+class CustomVimBuffer(vim_helper.VimBuffer):
+  @property
+  def filetypes(self):
+    return [filetype]
+
+vim_helper._orig_buf = vim_helper.buf
+vim_helper.buf = CustomVimBuffer()  # TODO: avoid creating a new class instance every time
+EOF
+endfunction
+
+function! cmp_nvim_ultisnips#reset_filetype()
+python3 << EOF
+from UltiSnips import vim_helper
+
+# Restore to original VimBuffer instance
+vim_helper.buf = vim_helper._orig_buf
+EOF
+endfunction
+
+function! cmp_nvim_ultisnips#setup_treesitter_autocmds()
+  augroup cmp_nvim_ultisnips
+    autocmd!
+    autocmd CursorMovedI * lua require("cmp_nvim_ultisnips.treesitter").set_filetype()
+    autocmd InsertLeave * lua require("cmp_nvim_ultisnips.treesitter").reset_filetype()
+  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 = "treesitter",
 }
 
 function M.get_user_config(config)
@@ -18,6 +19,13 @@ 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 == "treesitter" or arg == "ultisnips_default"
+      end,
+      "either 'treesitter' or 'ultisnips_default'",
+    },
   })
   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" then
+    vim.fn["cmp_nvim_ultisnips#setup_treesitter_autocmds"]()
+  end
   return self
 end
 

lua/cmp_nvim_ultisnips/treesitter.lua

@@ -0,0 +1,44 @@
+local M = {}
+
+-- Check if nvim-treesitter is available
+-- (there would be a lot of copied code without this dependency)
+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
+    return nil
+  end
+
+  local cur_node = ts_utils.get_node_at_cursor()
+  if cur_node then
+    local parser = ts_parsers.get_parser()
+    return parser:language_for_range({ cur_node:range() }):lang()
+  end
+  return nil
+end
+
+local cur_ft_at_cursor
+function M.set_filetype()
+  local new_ft = get_ft_at_cursor()
+  if new_ft ~= cur_ft_at_cursor and new_ft ~= vim.bo.filetype then
+    vim.fn["cmp_nvim_ultisnips#set_filetype"](new_ft)
+    cur_ft_at_cursor = new_ft
+  end
+end
+
+function M.reset_filetype()
+  if cur_ft_at_cursor ~= nil and cur_ft_at_cursor ~= vim.bo.filetype then
+    vim.fn["cmp_nvim_ultisnips#reset_filetype"]()
+    cur_ft_at_cursor = nil
+  end
+end
+
+return M