feat(treesitter): add metadata option for get_node_text

5 files changed, 55 insertions, 37 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 7c3f295560..b438595b5f 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -154,6 +154,9 @@ The following new APIs or features were added.
 • Treesitter captures can now be transformed by directives. This will allow
   more complicated dynamic language injections.
 
+• |vim.treesitter.query.get_node_text()| now accepts a `metadata` option for
+  writing custom directives using |vim.treesitter.query.add_directive()|.
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changes*
 
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 7bc9f3f9a7..693804497d 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -746,6 +746,9 @@ get_node_text({node}, {source}, {opts})
       • {opts}    (table|nil) Optional parameters.
                   • concat: (boolean) Concatenate result in a string (default
                     true)
+                  • metadata (table) Metadata of a specific capture. This
+                    would be set to `metadata[capture_id]` when using
+                    |vim.treesitter.query.add_directive()|.
 
     Return: ~
         (string[]|string|nil)
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 3e1bc5d1cb..8255c6f4fe 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -354,14 +354,6 @@ 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.
@@ -416,7 +408,7 @@ function LanguageTree:_get_injections()
 
         -- Lang should override any other language tag
         if name == 'language' and not lang then
-          lang = get_node_text(node, id, metadata, self._source) --[[@as string]]
+          lang = query.get_node_text(node, self._source, { metadata = metadata[id] })
         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
index 5ec8c67462..9136b596be 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -55,6 +55,38 @@ local function add_included_lang(base_langs, lang, ilang)
   return false
 end
 
+---@private
+---@param buf (number)
+---@param range (table)
+---@param concat (boolean)
+---@returns (string[]|string|nil)
+local function buf_range_get_text(buf, range, concat)
+  local lines
+  local start_row, start_col, end_row, end_col = unpack(range)
+  local eof_row = a.nvim_buf_line_count(buf)
+  if start_row >= eof_row then
+    return nil
+  end
+
+  if end_col == 0 then
+    lines = a.nvim_buf_get_lines(buf, start_row, end_row, true)
+    end_col = -1
+  else
+    lines = a.nvim_buf_get_lines(buf, start_row, end_row + 1, true)
+  end
+
+  if #lines > 0 then
+    if #lines == 1 then
+      lines[1] = string.sub(lines[1], start_col + 1, end_col)
+    else
+      lines[1] = string.sub(lines[1], start_col + 1)
+      lines[#lines] = string.sub(lines[#lines], 1, end_col)
+    end
+  end
+
+  return concat and table.concat(lines, '\n') or lines
+end
+
 --- Gets the list of files used to make up a query
 ---
 ---@param lang string Language to get query for
@@ -240,40 +272,22 @@ end
 ---@param source (number|string) Buffer or string from which the {node} is extracted
 ---@param opts (table|nil) Optional parameters.
 ---          - concat: (boolean) Concatenate result in a string (default true)
+---          - metadata (table) Metadata of a specific capture. This would be
+---            set to `metadata[capture_id]` when using
+---            |vim.treesitter.query.add_directive()|.
 ---@return (string[]|string|nil)
 function M.get_node_text(node, source, opts)
   opts = opts or {}
   local concat = vim.F.if_nil(opts.concat, true)
+  local metadata = opts.metadata or {}
 
-  local start_row, start_col, start_byte = node:start()
-  local end_row, end_col, end_byte = node:end_()
-
-  if type(source) == 'number' then
-    local eof_row = a.nvim_buf_line_count(source)
-    if start_row >= eof_row then
-      return nil
-    end
-
-    local lines ---@type string[]
-    if end_col == 0 then
-      lines = a.nvim_buf_get_lines(source, start_row, end_row, true)
-      end_col = -1
-    else
-      lines = a.nvim_buf_get_lines(source, start_row, end_row + 1, true)
-    end
-
-    if #lines > 0 then
-      if #lines == 1 then
-        lines[1] = string.sub(lines[1], start_col + 1, end_col)
-      else
-        lines[1] = string.sub(lines[1], start_col + 1)
-        lines[#lines] = string.sub(lines[#lines], 1, end_col)
-      end
-    end
-
-    return concat and table.concat(lines, '\n') or lines
+  if metadata.text then
+    return metadata.text
+  elseif type(source) == 'number' then
+    return metadata.range and buf_range_get_text(source, metadata.range, concat)
+      or buf_range_get_text(source, { node:range() }, concat)
   elseif type(source) == 'string' then
-    return source:sub(start_byte + 1, end_byte)
+    return source:sub(select(3, node:start()) + 1, select(3, node:end_()))
   end
 end
 
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index 0f187f3d52..2d0d57fbd0 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -277,6 +277,9 @@ end]]
       function fake_node:end_()
         return 3, 0, 23
       end
+      function fake_node:range()
+        return 3, 0, 3, 0
+      end
       return vim.treesitter.get_node_text(fake_node, 0) == nil
     ]])
     eq(true, result)
@@ -296,6 +299,9 @@ end]]
       function fake_node:end_()
         return 1, 0, 7
       end
+      function fake_node:range()
+        return 1, 0, 1, 0
+      end
       return vim.treesitter.get_node_text(fake_node, 0) == ''
     ]])
     eq(true, result)