Refactor treesitter filetype detection

autoload/vsnip/source.vim

@@ -23,12 +23,26 @@ endfunction
 "
 function! vsnip#source#filetypes( bufnr ) abort
   if has( "nvim" )
-    let l:filetype = v:lua.require'vsnip.treesitter'.get_ft_at_cursor( a:bufnr )
+    let l:filetypes = v:lua.require'vsnip.treesitter'.get_ft_at_cursor( a:bufnr )
+
+    " buffer has no filetype defined
+    if l:filetypes.filetype == ""
+      return [ "global" ]
+
+    " buffer has filetype
+    else
+      return
+        \ get( g:vsnip_filetypes, l:filetypes.injected_filetype,
+        \ get( g:vsnip_filetypes, l:filetypes.filetype,
+        \ [ l:filetypes.filetype ]
+        \ ) )
+        \ + [ "global" ]
+    endif
   else
     let l:filetype = getbufvar( a:bufnr, "&filetype", "" )
-  endif
 
-  return split( l:filetype, '\.' ) + get( g:vsnip_filetypes, l:filetype, [] ) + [ "global" ]
+    return split( l:filetype, '\.' ) + get( g:vsnip_filetypes, l:filetype, [] ) + [ "global" ]
+  endif
 endfunction
 
 "

doc/vsnip.txt

@@ -71,6 +71,16 @@ VARIABLE                                                      *vsnip-variable*
       let g:vsnip_filetypes = {}
       let g:vsnip_filetypes.javascriptreact = ['javascript']
 <
+
+    If you are using `treesitter` you can define snippets for injected
+    languages like this:
+
+>
+    let g:vsnip_filetypes['vim/lua'] = ['lua', 'vim/lua']
+    let g:vsnip_filetypes['vue'] = ['html']
+    let g:vsnip_filetypes['vue/javascript'] = ['javascript', 'vue/javascript']
+<
+
  let g:vsnip_deactivate_on = g:vsnip#DeactivateOn.OutsideOfSnippet~
     Specify when to deactivate the current snippet.
 

lua/vsnip/treesitter.lua

@@ -10,25 +10,56 @@ if not ok_utils then
   ts_utils = nil
 end
 
-function M.is_available ()
+local function get_parser_filetype ( lang )
+  if lang and ts_parsers.list[ lang ] then
+    return ts_parsers.list[ lang ].filetype or lang
+  else
+    return ""
+  end
+end
+
+local function is_available ()
   return ok_parsers and ok_utils
 end
 
 function M.get_ft_at_cursor ( bufnr )
-  if M.is_available() then
+  local filetypes = {
+    filetype = "",
+    injected_filetype = "",
+  }
+
+  if is_available() then
     local cur_node = ts_utils.get_node_at_cursor( vim.fn.bufwinid( bufnr ) )
 
     if cur_node then
       local parser = ts_parsers.get_parser( bufnr )
-      local lang = parser:language_for_range( { cur_node:range() } ):lang()
+      local language_tree_at_cursor = parser:language_for_range( { cur_node:range() } )
+      local language_at_cursor = language_tree_at_cursor:lang()
+
+      local filetype = get_parser_filetype( language_at_cursor )
 
-      if ts_parsers.list[ lang ] ~= nil then
-        return ts_parsers.list[ lang ].filetype or lang
+      if filetype ~= "" then
+        filetypes.filetype = filetype
+
+        local parent_language_tree = language_tree_at_cursor:parent()
+
+        if parent_language_tree then
+          local parent_language = parent_language_tree:lang()
+          local parent_filetype = get_parser_filetype( parent_language )
+
+          if parent_filetype ~= "" then
+            filetypes.injected_filetype = parent_filetype .. "/" .. filetype
+          end
+        end
+
+        return filetypes
       end
     end
   end
 
-  return vim.bo[ bufnr ].filetype or ""
+  filetypes.filetype = vim.bo[ bufnr ].filetype or ""
+
+  return filetypes
 end
 
 return M