neovim

commit bb8845340b1b9c2180fb19f049ff9deff5857d99
parent 5c4b503d3cb4a48d083bcf50d4932927e6eb749d
Author: figsoda <figsoda@pm.me>
Date:   Thu, 21 Jul 2022 12:08:37 +0100

feat(treesitter): allow capture text to be transformed

Co-authored-by: Lewis Russell <lewis6991@gmail.com>

Diffstat:
M
runtime/doc/news.txt
 | 
3
+++
M
runtime/lua/vim/treesitter/languagetree.lua
 | 
10
+++++++++-
M
runtime/lua/vim/treesitter/query.lua
 | 
15
+++++++++++++++
M
test/functional/treesitter/parser_spec.lua
 | 
2
+-

4 files changed, 28 insertions(+), 2 deletions(-)
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
@@ -151,6 +151,9 @@ The following new APIs or features were added.
 
 • |:highlight| now supports an additional attribute "altfont".
 
+• Treesitter captures can now be transformed by directives. This will allow
+  more complicated dynamic language injections.
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changes*
 
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
@@ -354,6 +354,14 @@ local function get_range_from_metadata(node, id, metadata)
   return { node:range() }
 end
 
+---@private
+local function get_node_text(node, id, metadata, source)
+  if metadata[id] and metadata[id].text then
+    return metadata[id].text
+  end
+  return query.get_node_text(node, source)
+end
+
 --- Gets language injection points by language.
 ---
 --- This is where most of the injection processing occurs.
@@ -408,7 +416,7 @@ function LanguageTree:_get_injections()
 
         -- Lang should override any other language tag
         if name == 'language' and not lang then
-          lang = query.get_node_text(node, self._source) --[[@as string]]
+          lang = get_node_text(node, id, metadata, self._source) --[[@as string]]
         elseif name == 'combined' then
           combined = true
         elseif name == 'content' and #ranges == 0 then
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
@@ -452,6 +452,21 @@ local directive_handlers = {
       metadata[capture_id].range = range
     end
   end,
+
+  -- Transform the content of the node
+  -- Example: (#gsub! @_node ".*%.(.*)" "%1")
+  ['gsub!'] = function(match, _, bufnr, pred, metadata)
+    assert(#pred == 4)
+
+    local id = pred[2]
+    local node = match[id]
+    local text = M.get_node_text(node, bufnr, { metadata = metadata[id] }) or ''
+
+    if not metadata[id] then
+      metadata[id] = {}
+    end
+    metadata[id].text = text:gsub(pred[3], pred[4])
+  end,
 }
 
 --- Adds a new predicate to be used in queries
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
@@ -728,7 +728,7 @@ int x = INT_MAX;
         return list
         ]]
 
-        eq({ 'offset!', 'set!' }, res_list)
+        eq({ 'gsub!', 'offset!', 'set!' }, res_list)
       end)
     end)
   end)