From dad48dad48e7964efab437d89624c2195ef966f9 Mon Sep 17 00:00:00 2001
From: danymat <d.danymat@gmail.com>
Date: Thu, 13 Jan 2022 11:54:18 +0100
Subject: [PATCH] feat: Allow recursive extractions

---
 lua/neogen/utilities/extractors.lua | 14 ++++++++------
 lua/neogen/utilities/nodes.lua      |  8 +++++++-
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/lua/neogen/utilities/extractors.lua b/lua/neogen/utilities/extractors.lua
index 6fd56a6..17ce1ef 100644
--- a/lua/neogen/utilities/extractors.lua
+++ b/lua/neogen/utilities/extractors.lua
@@ -2,11 +2,11 @@ local ts_utils = require("nvim-treesitter.ts_utils")
 
 neogen.utilities.extractors = {
     --- Extract the content from each node from data
-    --- @param _ any self
+    --- @param self any
     --- @param data table a list of k,v values where k is the node_type and v a table of nodes
     --- @param opts table|nil optional configurations (supported: type = true will get node's type instead of node's text)
     --- @return any result the same table as data but with node texts instead
-    extract_from_matched = function(_, data, opts)
+    extract_from_matched = function(self, data, opts)
         opts = opts or {}
         local result = {}
         for k, v in pairs(data) do
@@ -16,10 +16,12 @@ neogen.utilities.extractors = {
             local get_text = function(node)
                 return ts_utils.get_node_text(node)[1]
             end
-            if opts.type == true then
-                result[k] = vim.tbl_map(get_type, v)
-            else
-                result[k] = vim.tbl_map(get_text, v)
+
+            if type(v) == "table" then
+                result[k] = self:extract_from_matched(v, opts)
+            elseif type(v) == "userdata" then
+                local cb = opts.type and get_type or get_text
+                result[k] = cb(v)
             end
         end
         return result
diff --git a/lua/neogen/utilities/nodes.lua b/lua/neogen/utilities/nodes.lua
index e35abd9..3a4301a 100644
--- a/lua/neogen/utilities/nodes.lua
+++ b/lua/neogen/utilities/nodes.lua
@@ -91,7 +91,13 @@ neogen.utilities.nodes = {
                     if result[name] == nil then
                         result[name] = {}
                     end
-                    table.insert(result[name], child)
+                    local res = child
+                    if subtree.subtree then
+                        res = { res }
+                        local nodes = self:matching_nodes_from(child, subtree.subtree)
+                        table.insert(res, nodes)
+                    end
+                    table.insert(result[name], res)
                 else
                     local nodes = self:matching_nodes_from(child, subtree.subtree, result)
                     result = vim.tbl_deep_extend("keep", result, nodes)