Refactor to a simpler tree parsing

You can now parse the syntax tree more efficiently with the new
utilities functions.
For examples, please check out the lua and python configurations
This commit is contained in:
Daniel Mathiot
2021-08-24 09:41:15 +02:00
parent 790c98edf1
commit 77515a574d
8 changed files with 161 additions and 164 deletions

View File

@@ -197,26 +197,28 @@ end
```lua ```lua
data = { data = {
-- If function or local_function is found as a parent
["function|local_function"] = { ["function|local_function"] = {
-- Get second child from the parent node -- get second child from the parent node
["2"] = { ["2"] = {
-- This second child has to be of type "parameters", otherwise does nothing -- it has to be of type "parameters"
match = "parameters", match = "parameters",
-- Extractor function that returns a set of TSname = values with values being of type string[]
extract = function(node) extract = function(node)
local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(node) local tree = {
local varargs = neogen.utilities.extractors:extract_children_text("spread")(node) { retrieve = "all", node_type = "identifier", extract = true },
{ retrieve = "all", node_type = "spread", extract = true }
}
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
return { return {
parameters = regular_params, parameters = res.identifier,
vararg = varargs, vararg = res.spread,
} }
end, end,
}, },
}, },
}, }
``` ```
Notes: Notes:

View File

@@ -3,7 +3,6 @@ assert(ok, "neogen requires nvim-treesitter to operate :(")
neogen = {} neogen = {}
-- Require utilities -- Require utilities
neogen.utilities = {} neogen.utilities = {}
require("neogen.utilities.extractors") require("neogen.utilities.extractors")

View File

@@ -1,5 +1,3 @@
local ts_utils = require("nvim-treesitter.ts_utils")
return { return {
-- Search for these nodes -- Search for these nodes
parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" }, parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" },
@@ -13,12 +11,16 @@ return {
match = "parameters", match = "parameters",
extract = function(node) extract = function(node)
local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(node) local tree = {
local varargs = neogen.utilities.extractors:extract_children_text("spread")(node) { retrieve = "all", node_type = "identifier", extract = true },
{ retrieve = "all", node_type = "spread", extract = true },
}
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
return { return {
parameters = regular_params, parameters = res.identifier,
vararg = varargs, vararg = res.spread,
} }
end, end,
}, },
@@ -28,20 +30,25 @@ return {
match = "function_definition", match = "function_definition",
extract = function(node) extract = function(node)
local regular_params = neogen.utilities.extractors:extract_children_from({ local tree = {
[1] = "extract", {
}, "identifier")(node) retrieve = "first",
node_type = "parameters",
subtree = {
{ retrieve = "all", node_type = "identifier", extract = true },
{ retrieve = "all", node_type = "spread", extract = true },
},
},
{ retrieve = "first", node_type = "return_statement", extract = true },
}
local varargs = neogen.utilities.extractors:extract_children_from({ local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
[1] = "extract", local res = neogen.utilities.extractors:extract_from_matched(nodes)
}, "spread")(node)
local return_statement = neogen.utilities.extractors:extract_children_text("return_statement")(node)
return { return {
parameters = regular_params, parameters = res.identifier,
vararg = varargs, vararg = res.spread,
return_statement = return_statement, return_statement = res.return_statement,
} }
end, end,
}, },

View File

@@ -9,67 +9,47 @@ return {
["function_definition"] = { ["function_definition"] = {
["0"] = { ["0"] = {
extract = function(node) extract = function(node)
local results = { local results = {}
parameters = {},
return_statement = {} local tree = {
{
retrieve = "all",
node_type = "parameters",
subtree = {
{ retrieve = "all", node_type = "identifier", extract = true },
{
retrieve = "all",
node_type = "default_parameter",
subtree = { { retrieve = "all", node_type = "identifier", extract = true } },
},
{
retrieve = "all",
node_type = "typed_parameter",
subtree = { { retrieve = "all", node_type = "identifier", extract = true } },
},
{
retrieve = "all",
node_type = "typed_default_parameter",
subtree = { { retrieve = "all", node_type = "identifier", extract = true } },
},
},
},
{
retrieve = "first",
node_type = "block",
subtree = {
{ retrieve = "all", node_type = "return_statement", extract = true },
},
},
} }
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
local params = neogen.utilities.nodes:matching_child_nodes(node, "parameters")[1] results.parameters = res.identifier
results.return_statement = res.return_statement
if #params == 0 then
results.parameters = nil
end
local found_nodes
-- Find regular parameters
local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(params)
if #regular_params == 0 then
regular_params = nil
end
for _, _params in pairs(regular_params) do
table.insert(results.parameters, _params)
end
results.parameters = regular_params
-- Find regular optional parameters
found_nodes = neogen.utilities.nodes:matching_child_nodes(params, "default_parameter")
for _,_node in pairs(found_nodes) do
local _params = neogen.utilities.extractors:extract_children_text("identifier")(_node)[1]
table.insert(results.parameters, _params)
end
-- Find typed params
found_nodes = neogen.utilities.nodes:matching_child_nodes(params, "typed_parameter")
for _,_node in pairs(found_nodes) do
local _params = neogen.utilities.extractors:extract_children_text("identifier")(_node)[1]
table.insert(results.parameters, _params)
end
-- TODO Find optional typed params
found_nodes = neogen.utilities.nodes:matching_child_nodes(params, "typed_default_parameter")
for _,_node in pairs(found_nodes) do
local _params = neogen.utilities.extractors:extract_children_text("identifier")(_node)[1]
table.insert(results.parameters, _params)
end
local body = neogen.utilities.nodes:matching_child_nodes(node, "block")[1]
if body ~= nil then
local return_statement = neogen.utilities.nodes:matching_child_nodes(body, "return_statement")
if #return_statement == 0 then
return_statement = nil
end
results.return_statement = return_statement
end
return results return results
end end,
}, },
}, },
["class_definition"] = { ["class_definition"] = {
@@ -77,34 +57,40 @@ return {
match = "block", match = "block",
extract = function(node) extract = function(node)
local results = { local results = {}
attributes = {} local tree = {
{
retrieve = "first",
node_type = "function_definition",
subtree = {
{
retrieve = "first",
node_type = "block",
subtree = {
{
retrieve = "all",
node_type = "expression_statement",
subtree = {
{ retrieve = "first", node_type = "assignment", extract = true },
},
},
},
},
},
},
} }
local init_function = neogen.utilities.nodes:matching_child_nodes(node, "function_definition")[1] local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
if init_function == nil then results.attributes = {}
return for _, assignment in pairs(nodes["assignment"]) do
end
local body = neogen.utilities.nodes:matching_child_nodes(init_function, "block")[1]
if body == nil then
return
end
local expressions = neogen.utilities.nodes:matching_child_nodes(body, "expression_statement")
for _,expression in pairs(expressions) do
local assignment = neogen.utilities.nodes:matching_child_nodes(expression, "assignment")[1]
if assignment ~= nil then
local left_side = assignment:field("left")[1] local left_side = assignment:field("left")[1]
local left_attribute = left_side:field("attribute")[1] local left_attribute = left_side:field("attribute")[1]
table.insert(results.attributes, ts_utils.get_node_text(left_attribute)[1]) table.insert(results.attributes, ts_utils.get_node_text(left_attribute)[1])
end end
end
return results return results
end end,
}, },
}, },
}, },
@@ -115,7 +101,7 @@ return {
generator = nil, generator = nil,
template = { template = {
annotation_convention = "numpydoc", -- required: Which annotation convention to use (default_generator) annotation_convention = "google_docstrings", -- required: Which annotation convention to use (default_generator)
append = { position = "after", child_name = "block" }, -- optional: where to append the text (default_generator) append = { position = "after", child_name = "block" }, -- optional: where to append the text (default_generator)
use_default_comment = false, -- If you want to prefix the template with the default comment for the language, e.g for python: # (default_generator) use_default_comment = false, -- If you want to prefix the template with the default comment for the language, e.g for python: # (default_generator)
google_docstrings = { google_docstrings = {
@@ -132,7 +118,7 @@ return {
{ "attributes", "%s: ", { before_first_item = { "", "Attributes", "----------" } } }, { "attributes", "%s: ", { before_first_item = { "", "Attributes", "----------" } } },
{ "return_statement", "", { before_first_item = { "", "Returns", "-------" } } }, { "return_statement", "", { before_first_item = { "", "Returns", "-------" } } },
{ nil, "" }, { nil, "" },
{ nil, '"""' } { nil, '"""' },
} },
}, },
} }

View File

@@ -1,50 +1,18 @@
local ts_utils = require("nvim-treesitter.ts_utils") local ts_utils = require("nvim-treesitter.ts_utils")
neogen.utilities.extractors = { neogen.utilities.extractors = {
--- Return a function to extract content of required children from a node --- Extract the content from each node from data
--- @param _ any self --- @param _ any self
--- @param name string the children we want to extract (if multiple childrens, separate each one with "|") --- @param data table a list of k,v values where k is the node_type and v a table of nodes
--- @return function cb function taking a node and getting the content of each children we want from name --- @return any result the same table as data but with node texts instead
extract_children_text = function(_, name) extract_from_matched = function(_, data)
return function(node)
local result = {} local result = {}
local split = vim.split(name, "|", true) for k, v in pairs(data) do
local get_text = function(node)
for child in node:iter_children() do return ts_utils.get_node_text(node)[1]
if vim.tbl_contains(split, child:type()) then
table.insert(result, ts_utils.get_node_text(child)[1])
end end
result[k] = vim.tbl_map(get_text, v)
end end
return result return result
end
end,
--- Extract content from specified children from a tree
--- the tree parameter can be a nested { [key] = value} with key being the
--- * key: is which children we want to extract the values from (e.g first children is 1)
--- * value: "extract" or { [key] = value }. If value is "extract", it will extract the key child node
--- Example (extract the first child node from the first child node of the parent node):
--- [1] = {
--- [1] = "extract"
--- }
--- @param tree table see description
--- @param name string the children we want to extract (if multiple children, separate each one with "|")
extract_children_from = function(self, tree, name)
return function(node)
local result = {}
for i, subtree in pairs(tree) do
local child_node = node:named_child(tonumber(i) - 1)
if subtree == "extract" then
return self:extract_children_text(name)(child_node)
else
return self:extract_children_from(subtree, name)(child_node)
end
end
return result
end
end, end,
} }

View File

@@ -2,16 +2,51 @@ neogen.utilities.nodes = {
--- Get a list of child nodes that match the provided node name --- Get a list of child nodes that match the provided node name
--- @param _ any --- @param _ any
--- @param parent userdata the parent's node --- @param parent userdata the parent's node
--- @param node_name string the node type to search for --- @param node_name string the node type to search for (if multiple childrens, separate each one with "|")
--- @return table a table of nodes that matched the name --- @return table a table of nodes that matched the name
matching_child_nodes = function(_, parent, node_name) matching_child_nodes = function(_, parent, node_name)
local results = {} local results = {}
local split = vim.split(node_name, "|", true)
for child in parent:iter_children() do for child in parent:iter_children() do
if child:type() == node_name then if vim.tbl_contains(split, child:type()) then
table.insert(results, child) table.insert(results, child)
end end
end end
return results return results
end, end,
--- Get all required nodes from tree
--- @param parent userdata the parent node
--- @param tree table a nested table : { retrieve = "all|first", node_type = node_name, subtree = tree }
--- If you want to extract the node, do not specify the subtree and instead: extract = true
--- @param result table the table of results
--- @return table result a table of k,v where k are node_types and v all matched nodes
matching_nodes_from = function(self, parent, tree, result)
result = result or {}
for _, subtree in pairs(tree) do
-- Match all child nodes
local matched = self:matching_child_nodes(parent, subtree.node_type)
-- Only keep first matched child node
if subtree.retrieve == "first" and #matched ~= 0 then
matched = { matched[1] }
end
for _, child in pairs(matched) do
-- Add to results
if subtree.extract == true then
if result[subtree.node_type] == nil then
result[subtree.node_type] = {}
end
table.insert(result[subtree.node_type], child)
else
local test = self:matching_nodes_from(child, subtree.subtree, result)
result = vim.tbl_deep_extend("keep", result, test)
end
end
end
return result
end,
} }