Better code chunk extraction

lua/r/maps.lua

@@ -294,8 +294,8 @@ local send = function(file_type)
     end
     if file_type == "rmd" or file_type == "quarto" then
         create_maps("nvi", "RKnit",           "kn", "<Cmd>lua require('r.run').knit()")
-        create_maps("ni",  "RSendChunk",      "cc", "<Cmd>lua require('r.rmd').send_R_chunk(false)")
-        create_maps("ni",  "RDSendChunk",     "cd", "<Cmd>lua require('r.rmd').send_R_chunk(true)")
+        create_maps("ni",  "RSendChunk",      "cc", "<Cmd>lua require('r.rmd').send_current_chunk(false)")
+        create_maps("ni",  "RDSendChunk",     "cd", "<Cmd>lua require('r.rmd').send_current_chunk(true)")
         create_maps("n",   "RNextRChunk",     "gn", "<Cmd>lua require('r.rmd').next_chunk()")
         create_maps("n",   "RPreviousRChunk", "gN", "<Cmd>lua require('r.rmd').previous_chunk()")
     end

lua/r/quarto.lua

@@ -1,5 +1,33 @@
 local M = {}
 
+local Chunk = {}
+Chunk.__index = Chunk
+
+--- Constructor for the Chunk class
+---@param content string The content of the code chunk.
+---@param start_row integer The starting row of the code chunk.
+---@param end_row integer The ending row of the code chunk.
+---@param info_string_params table The parameters specified in the info string of the code chunk.
+---@param comment_params table The parameters specified in the code chunk with #|
+---@param lang string The language of the code chunk.
+---@return table
+function Chunk:new(content, start_row, end_row, info_string_params, comment_params, lang)
+    local chunk = {
+        content = content,
+        start_row = start_row,
+        end_row = end_row,
+        info_string_params = info_string_params,
+        comment_params = comment_params,
+        lang = lang,
+    }
+
+    setmetatable(chunk, Chunk)
+
+    return chunk
+end
+
+function Chunk:range() return self.start_row, self.end_row end
+
 M.command = function(what)
     local config = require("r.config").get_config()
     local send_cmd = require("r.send").cmd
@@ -23,43 +51,217 @@ end
 
 --- Helper function to get code block from Rmd or Quarto document
 ---@param bufnr  integer The buffer number.
----@param lang string The language of the code chunk.
----@param row integer|nil The row number. If nil, all code chunks are returned.
 ---@return table|nil
-M.get_code_chunk = function(bufnr, lang, row)
+local get_code_chunks = function(bufnr)
     local root = require("r.utils").get_root_node()
-
     if not root then return nil end
 
     local query = vim.treesitter.query.parse(
         "markdown",
-        string.format(
-            [[
+        [[
                 (fenced_code_block
-                (info_string (language) @lang (#eq? @lang "%s"))
-                (code_fence_content) @content)
-            ]],
-            lang
-        )
+                    (info_string (language) @lang) @info_string
+                    (#match? @info_string "^\\{.*\\}$")
+                    (code_fence_content) @content) @fenced_code_block
+            ]]
     )
 
     bufnr = bufnr or vim.api.nvim_get_current_buf()
+    local code_chunks = {}
 
-    local r_contents = {}
-    for _, node, _ in query:iter_captures(root, bufnr, 0, -1) do
-        if node:type() == "code_fence_content" then
+    for id, node, _ in query:iter_captures(root, bufnr, 0, -1) do
+        local capture_name = query.captures[id]
+
+        if capture_name == "content" then
+            local lang
+            local info_string_params = {}
             local start_row, _, end_row, _ = node:range()
-            if not row or (row >= start_row and row <= end_row) then
-                table.insert(r_contents, {
-                    content = vim.treesitter.get_node_text(node, bufnr),
-                    start_row = start_row,
-                    end_row = end_row,
-                })
-                if row then break end
+
+            -- Get the info string of the code block and parse it
+            local parent = node:parent()
+            if parent then
+                local info_node = parent:child(1)
+                if info_node and info_node:type() == "info_string" then
+                    local info_string = vim.treesitter.get_node_text(info_node, bufnr)
+                    lang, info_string_params = M.parse_info_string_params(info_string)
+                end
             end
+
+            -- Get the parameters specified in the code chunk with #|
+            local comment_params =
+                M.parse_comment_params(vim.treesitter.get_node_text(node, bufnr))
+
+            local chunk = Chunk:new(
+                vim.treesitter.get_node_text(node, bufnr),
+                start_row,
+                end_row,
+                info_string_params,
+                comment_params,
+                lang
+            )
+
+            table.insert(code_chunks, chunk)
         end
     end
-    return r_contents
+
+    return code_chunks
+end
+
+--- Helper function to parse the info string of a code block
+---@param info_string string The info string of the code block.
+---@return string,table
+M.parse_info_string_params = function(info_string)
+    local params = {}
+
+    if info_string == nil then return "", params end
+
+    local lang = info_string:match("^%s*{?([^%s,{}]+)")
+    if lang == nil then return "", params end
+
+    local param_str = info_string:sub(#lang + 1):gsub("[{}]", "") -- Remove { and }
+    local param_list = vim.split(param_str, ",")
+    for _, param in ipairs(param_list) do
+        local key, value = param:match("^%s*([^=]+)%s*=%s*([^%s]+)%s*$")
+
+        if key and value then
+            key = key:match("^%s*(.-)%s*$") -- Trim leading/trailing spaces from key
+            value = value:match("^%s*(.-)%s*$") -- Trim leading/trailing spaces from value
+            params[key] = value
+        end
+    end
+
+    return lang, params
+end
+
+--- Helper function to parse the parameters specified in the code chunk with #|
+---@param code_content string The content of the code chunk.
+---@return table
+M.parse_comment_params = function(code_content)
+    local params = {}
+
+    for line in code_content:gmatch("[^\r\n]+") do
+        local key, value = line:match("^#|%s*([^:]+)%s*:%s*(.+)%s*$")
+        if key and value then
+            params[key:match("^%s*(.-)%s*$")] = value:match("^%s*(.-)%s*$")
+        end
+    end
+
+    return params
+end
+
+--- This function gets the current code chunk based on the cursor position
+---@param bufnr integer The buffer number.
+---@return table
+M.get_current_code_chunk = function(bufnr)
+    local row, _ = unpack(vim.api.nvim_win_get_cursor(0))
+
+    local chunks = get_code_chunks(bufnr)
+    if not chunks then return {} end
+
+    for _, chunk in ipairs(chunks) do
+        local chunk_start_row, chunk_end_row = chunk:range()
+        if row >= chunk_start_row and row <= chunk_end_row then return chunk end
+    end
+
+    return {}
+end
+
+-- Function to get all code chunks above the current cursor position
+---@param bufnr integer The buffer number.
+---@return table
+M.get_chunks_above_cursor = function(bufnr)
+    local row, _ = unpack(vim.api.nvim_win_get_cursor(0))
+
+    local chunks = get_code_chunks(bufnr)
+
+    if not chunks then return {} end
+
+    local chunks_above = {}
+
+    for _, chunk in ipairs(chunks) do
+        local _, chunk_end_row = chunk:range()
+
+        if chunk_end_row < row then table.insert(chunks_above, chunk) end
+    end
+
+    return chunks_above
+end
+
+--- This function gets all the code chunks in the buffer
+---@param bufnr integer The buffer number.
+---@return table|nil
+M.get_all_code_chunks = function(bufnr)
+    local chunks = get_code_chunks(bufnr)
+    return chunks
+end
+
+--- This function filters the code chunks based on the language
+---@param chunks table The code chunks.
+---@param langs string[] The languages to filter by.
+---@return table The filtered code chunks.
+M.filter_code_chunks_by_lang = function(chunks, langs)
+    local lang_set = {}
+    for _, lang in ipairs(langs) do
+        lang_set[lang] = true
+    end
+
+    return vim.tbl_filter(
+        function(chunk) return type(chunk) == "table" and lang_set[chunk.lang] or false end,
+        chunks
+    )
+end
+
+--- This function filters the code chunks based on the eval parameter. If the eval parameter is not found it is assumed to be true
+---@param chunks table
+---@return table
+M.filter_code_chunks_by_eval = function(chunks)
+    -- If chunks is a single chunk (table), wrap it in a table to ensure uniform processing
+    if type(chunks) ~= "table" or (type(chunks) == "table" and #chunks == 0) then
+        chunks = { chunks }
+    end
+
+    return vim.tbl_filter(function(chunk)
+        if type(chunk) ~= "table" then
+            return false -- Skip this chunk if it’s not a table
+        end
+
+        -- Default eval is true if not provided
+        local eval = true
+
+        -- Check for eval in comment_params
+        if chunk.comment_params and chunk.comment_params.eval then
+            eval = chunk.comment_params.eval == "true"
+        -- Check for eval in info_string_params
+        elseif chunk.info_string_params and chunk.info_string_params.eval then
+            eval = chunk.info_string_params.eval == "TRUE"
+        end
+
+        return eval -- Return true if eval is "true"
+    end, chunks)
+end
+
+--- Formats the code chunks into a list of code lines that can be executed in
+--- R. The code lines are formatted based on the language of the code chunk. If
+--- the language is python, the code line is wrapped in
+--- reticulate::py_run_string. If the language is R, the code line is returned
+--- as is.
+---@param chunks table The code chunks.
+---@return table
+M.codelines_from_chunks = function(chunks)
+    local codelines = {}
+
+    for _, chunk in ipairs(chunks) do
+        local lang = chunk.lang
+        local content = chunk.content
+
+        if lang == "python" then
+            table.insert(codelines, 'reticulate::py_run_string("' .. content .. '")')
+        elseif lang == "r" then
+            table.insert(codelines, content)
+        end
+    end
+
+    return codelines
 end
 
 return M

lua/r/rmd.lua

@@ -1,9 +1,10 @@
 local inform = require("r.log").inform
 local config = require("r.config").get_config()
-local send = require("r.send")
 local get_lang = require("r.utils").get_lang
 local uv = vim.uv
 local chunk_key = nil
+local quarto = require("r.quarto")
+local source_chunk = require("r.send").source_chunk
 
 local M = {}
 
@@ -49,44 +50,25 @@ M.write_chunk = function()
     end
 end
 
--- Internal function to send a Python code chunk to R for execution.
--- This is not exposed in the module table `M` and is only called within `M.send_R_chunk`.
----@param m boolean If true, moves to the next chunk after sending the current one.
-local send_py_chunk = function(m)
-    -- Find the start and end of Python code chunk
-    local chunkline = vim.fn.search("^[ \t]*```[ ]*{python", "bncW") + 1
-    local docline = vim.fn.search("^[ \t]*```", "ncW") - 1
-    local lines = vim.api.nvim_buf_get_lines(0, chunkline - 1, docline, true) -- Get chunk lines
-    local ok = send.source_lines(lines, "PythonCode")
-    if ok == 0 then return end -- check if sending was successful
-    if m == true then M.next_chunk() end -- optional: move to next chunk
-end
+--- Sends the current R or Python code chunk to the R console for evaluation.
+---@param m boolean If true, the cursor will move to the next code chunk after evaluation.
+M.send_current_chunk = function(m)
+    local bufnr = vim.api.nvim_get_current_buf()
 
---- Sends the current R code chunk to R for execution.
--- This function ensures the cursor is positioned inside an R code chunk before attempting to send it.
--- If inside a Python code chunk, it will delegate to `send_py_chunk`.
----@param m boolean If true, moves to the next chunk after sending the current one.
-M.send_R_chunk = function(m)
-    -- Ensure cursor is at the start of an R code chunk
-    if vim.api.nvim_get_current_line():find("^%s*```%s*{r") then
-        local lnum = vim.api.nvim_win_get_cursor(0)[1]
-        vim.api.nvim_win_set_cursor(0, { lnum + 1, 0 })
-    end
-    -- Check for R code chunk; if not, check for Python code chunk
-    local lang = get_lang()
-    if lang ~= "r" then
-        if lang ~= "python" then
-            inform("Not inside an R code chunk.")
-        else
-            send_py_chunk(m)
-        end
+    local chunks = quarto.get_current_code_chunk(bufnr)
+    chunks = quarto.filter_code_chunks_by_eval(chunks)
+    chunks = quarto.filter_code_chunks_by_lang(chunks, { "r", "python" })
+
+    if #chunks == 0 then
+        inform("No R or Python code chunk found at the cursor position.")
         return
     end
-    -- find and send R chunk for execution
-    local chunkline = vim.fn.search("^[ \t]*```[ ]*{r", "bncW") + 1
-    local docline = vim.fn.search("^[ \t]*```", "ncW") - 1
-    local lines = vim.api.nvim_buf_get_lines(0, chunkline - 1, docline, true)
-    local ok = send.source_lines(lines, nil)
+
+    local codelines = quarto.codelines_from_chunks(chunks)
+
+    local lines = table.concat(codelines, "\n")
+
+    local ok = source_chunk(lines)
     if ok == 0 then return end
     if m == true then M.next_chunk() end
 end