feat: Allow recursive extractions

This commit is contained in:
danymat
2022-01-13 11:54:18 +01:00
parent a79fa586a9
commit dad48dad48
2 changed files with 15 additions and 7 deletions

View File

@@ -2,11 +2,11 @@ local ts_utils = require("nvim-treesitter.ts_utils")
neogen.utilities.extractors = { neogen.utilities.extractors = {
--- Extract the content from each node from data --- 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 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) --- @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 --- @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 {} opts = opts or {}
local result = {} local result = {}
for k, v in pairs(data) do for k, v in pairs(data) do
@@ -16,10 +16,12 @@ neogen.utilities.extractors = {
local get_text = function(node) local get_text = function(node)
return ts_utils.get_node_text(node)[1] return ts_utils.get_node_text(node)[1]
end end
if opts.type == true then
result[k] = vim.tbl_map(get_type, v) if type(v) == "table" then
else result[k] = self:extract_from_matched(v, opts)
result[k] = vim.tbl_map(get_text, v) elseif type(v) == "userdata" then
local cb = opts.type and get_type or get_text
result[k] = cb(v)
end end
end end
return result return result

View File

@@ -91,7 +91,13 @@ neogen.utilities.nodes = {
if result[name] == nil then if result[name] == nil then
result[name] = {} result[name] = {}
end 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 else
local nodes = self:matching_nodes_from(child, subtree.subtree, result) local nodes = self:matching_nodes_from(child, subtree.subtree, result)
result = vim.tbl_deep_extend("keep", result, nodes) result = vim.tbl_deep_extend("keep", result, nodes)