feat(ts): add support for multiline nodes in get_node_text (#14999)

mfussenegger · web-flow · commit 1f3c0593eb1d · 2021-12-12T12:05:39.000+01:00
Based on #14445 This extends `vim.treesitter.query.get_node_text` to return the text that spans a node's range even if start_row ~= end_row.
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
@@ -164,8 +164,6 @@ function M.parse_query(lang, query)
   return self
 end
 
--- TODO(vigoux): support multiline nodes too
-
 --- Gets the text corresponding to a given node
 ---
 ---@param node the node
@@ -175,11 +173,26 @@ function M.get_node_text(node, source)
   local end_row, end_col, end_byte = node:end_()
 
   if type(source) == "number" then
-    if start_row ~= end_row then
+    local lines
+    local eof_row = vim.api.nvim_buf_line_count(source)
+    if start_row >= eof_row then
       return nil
     end
-    local line = a.nvim_buf_get_lines(source, start_row, start_row+1, true)[1]
-    return string.sub(line, start_col+1, end_col)
+    if end_col == 0 then
+      lines = a.nvim_buf_get_lines(source, start_row, end_row, true)
+      end_col = #lines[#lines]
+    else
+      lines = a.nvim_buf_get_lines(source, start_row, end_row + 1, true)
+    end
+    lines[1] = string.sub(lines[1], start_col + 1)
+
+    local end_index = end_col
+    if #lines == 1 then
+      end_index = end_col - start_col
+    end
+    lines[#lines] = string.sub(lines[#lines], 1, end_index)
+
+    return table.concat(lines, "\n")
   elseif type(source) == "string" then
     return source:sub(start_byte+1, end_byte)
   end
@@ -211,11 +224,6 @@ local predicate_handlers = {
   ["lua-match?"] = function(match, _, source, predicate)
       local node = match[predicate[2]]
       local regex = predicate[3]
-      local start_row, _, end_row, _ = node:range()
-      if start_row ~= end_row then
-        return false
-      end
-
       return string.find(M.get_node_text(node, source), regex)
   end,
 
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
@@ -227,6 +227,43 @@ void ui_refresh(void)
     }, res)
   end)
 
+  it('supports getting text of multiline node', function()
+    if pending_c_parser(pending) then return end
+    insert(test_text)
+    local res = exec_lua([[
+      local parser = vim.treesitter.get_parser(0, "c")
+      local tree = parser:parse()[1]
+      return vim.treesitter.get_node_text(tree:root(), 0)
+    ]])
+    eq(test_text, res)
+
+    local res2 = exec_lua([[
+      local parser = vim.treesitter.get_parser(0, "c")
+      local root = parser:parse()[1]:root()
+      return vim.treesitter.get_node_text(root:child(0):child(0), 0)
+    ]])
+    eq('void', res2)
+  end)
+
+  it('support getting text where start of node is past EOF', function()
+    local text = [[
+def run
+  a = <<~E
+end]]
+    insert(text)
+    local result = exec_lua([[
+      local fake_node = {}
+      function fake_node:start()
+        return 3, 0, 23
+      end
+      function fake_node:end_()
+        return 3, 0, 23
+      end
+      return vim.treesitter.get_node_text(fake_node, 0) == nil
+    ]])
+    eq(true, result)
+  end)
+
   it('can match special regex characters like \\ * + ( with `vim-match?`', function()
     insert('char* astring = "\\n"; (1 + 1) * 2 != 2;')
 
