From ccbfa12586f3216495454b72860acdb2791dc150 Mon Sep 17 00:00:00 2001
From: Philippe Massicotte <pmassicotte@hotmail.com>
Date: Sun, 23 Feb 2025 14:05:38 -0500
Subject: [PATCH] refactor(numbers.lua, quarto.lua): generalize code chunk
 retrieval function

Replace `get_r_chunks_from_quarto` with `get_code_chunk` to support
retrieving code chunks of any specified language, not just R.
---
 lua/r/format/numbers.lua |  4 ++--
 lua/r/quarto.lua         | 15 +++++++++------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/lua/r/format/numbers.lua b/lua/r/format/numbers.lua
index e4a58cae..5147a90a 100644
--- a/lua/r/format/numbers.lua
+++ b/lua/r/format/numbers.lua
@@ -1,5 +1,5 @@
 local inform = require("r.log").inform
-local get_r_chunks_from_quarto = require("r.quarto").get_r_chunks_from_quarto
+local get_code_chunk = require("r.quarto").get_code_chunk
 local get_root_node = require("r.utils").get_root_node
 
 local M = {}
@@ -68,7 +68,7 @@ M.formatnum = function(bufnr)
     if not root_node then return end
 
     if filetype == "quarto" or filetype == "rmd" then
-        local r_chunks_content = get_r_chunks_from_quarto(root_node, bufnr)
+        local r_chunks_content = get_code_chunk(root_node, bufnr, nil, "r")
 
         for _, r_chunk in ipairs(r_chunks_content) do
             find_and_replace_float(r_chunk.content, bufnr, r_chunk.start_row, 0)
diff --git a/lua/r/quarto.lua b/lua/r/quarto.lua
index a9cb4fdd..c0423d76 100644
--- a/lua/r/quarto.lua
+++ b/lua/r/quarto.lua
@@ -26,14 +26,17 @@ end
 -- @param bufnr The buffer number
 -- @param cursor_pos The cursor position (optional)
 -- @return A table containing R code blocks with their content, start row and end row
-M.get_r_chunks_from_quarto = function(root, bufnr, cursor_pos)
+M.get_code_chunk = function(root, bufnr, cursor_pos, code_fence_language)
     local query = vim.treesitter.query.parse(
         "markdown",
-        [[
-        (fenced_code_block
-          (info_string (language) @lang (#eq? @lang "r"))
-          (code_fence_content) @content)
-        ]]
+        string.format(
+            [[
+                (fenced_code_block
+                (info_string (language) @lang (#eq? @lang "%s"))
+                (code_fence_content) @content)
+            ]],
+            code_fence_language
+        )
     )
 
     bufnr = bufnr or vim.api.nvim_get_current_buf()