Add support for multiple types, see Readme:Usage

This commit is contained in:
Daniel Mathiot
2021-08-25 13:39:55 +02:00
parent f5390e1c1f
commit 1002ae5948
5 changed files with 174 additions and 117 deletions

View File

@@ -13,7 +13,10 @@ require("neogen.locators.default")
require("neogen.granulators.default")
require("neogen.generators.default")
neogen.auto_generate = function(custom_template)
neogen.generate = function(opts)
opts = opts or {
type = "func",
}
vim.treesitter.get_parser(0):for_each_tree(function(tree, language_tree)
local language = neogen.configuration.languages[language_tree:lang()]
@@ -22,26 +25,26 @@ neogen.auto_generate = function(custom_template)
language.granulator = language.granulator or neogen.default_granulator
language.generator = language.generator or neogen.default_generator
if not language.parent[opts.type] or not language.data[opts.type] then
return
end
-- Use the language locator to locate one of the required parent nodes above the cursor
local located_parent_node = language.locator({
root = tree:root(),
current = ts_utils.get_node_at_cursor(0),
}, language.parent)
}, language.parent[opts.type])
if not located_parent_node then
return
end
-- Use the language granulator to get the required content inside the node found with the locator
local data = language.granulator(located_parent_node, language.data)
local data = language.granulator(located_parent_node, language.data[opts.type])
if data then
-- Will try to generate the documentation from a template and the data found from the granulator
local to_place, start_column, content = language.generator(
located_parent_node,
data,
custom_template or language.template
)
local to_place, start_column, content = language.generator(located_parent_node, data, language.template)
-- Append the annotation in required place
vim.fn.append(to_place, content)
@@ -57,7 +60,7 @@ neogen.auto_generate = function(custom_template)
end
function neogen.generate_command()
vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").auto_generate()')
vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").generate()')
end
neogen.setup = function(opts)

View File

@@ -15,38 +15,40 @@ local function_tree = {
},
}
return {
parent = { "function_declaration", "expression_statement", "variable_declaration" },
parent = { func = { "function_declaration", "expression_statement", "variable_declaration" }, },
data = {
["function_declaration"] = {
["0"] = {
func = {
["function_declaration"] = {
["0"] = {
extract = function(node)
local results = {}
local tree = function_tree
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
extract = function(node)
local results = {}
local tree = function_tree
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
results.parameters = res.identifier
results.return_statement = res.return_statement
return results
end,
results.parameters = res.identifier
results.return_statement = res.return_statement
return results
end,
},
},
},
["expression_statement|variable_declaration"] = {
["1"] = {
extract = function(node)
local results = {}
local tree = { { retrieve = "all", node_type = "function", subtree = function_tree } }
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
["expression_statement|variable_declaration"] = {
["1"] = {
extract = function(node)
local results = {}
local tree = { { retrieve = "all", node_type = "function", subtree = function_tree } }
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
results.parameters = res.identifier
results.return_statement = res.return_statement
return results
end,
results.parameters = res.identifier
results.return_statement = res.return_statement
return results
end,
},
},
},
}
},
template = {

View File

@@ -23,22 +23,49 @@ end
return {
-- Search for these nodes
parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" },
parent = {
func = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" },
class = { "local_variable_declaration" },
},
data = {
-- When the function is inside one of those
["local_variable_declaration|field|variable_declaration"] = {
["2"] = {
match = "function_definition",
func = {
-- When the function is inside one of those
["local_variable_declaration|field|variable_declaration"] = {
["2"] = {
match = "function_definition",
extract = common_function_extractor,
extract = common_function_extractor,
},
},
-- When the function is in the root tree
["function_definition|function|local_function"] = {
["0"] = {
extract = common_function_extractor,
},
},
},
-- When the function is in the root tree
["function_definition|function|local_function"] = {
["0"] = {
extract = common_function_extractor,
class = {
["local_variable_declaration"] = {
["0"] = {
extract = function(node)
local tree = {
{
retrieve = "first",
node_type = "variable_declarator",
subtree = {
{ retrieve = "first", node_type = "identifier", extract = true },
},
},
}
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
return {
class_name = res.identifier,
}
end,
},
},
},
},
@@ -59,6 +86,7 @@ return {
{ "parameters", "- @param %s any" },
{ "vararg", "- @vararg any" },
{ "return_statement", "- @return any" },
{ "class_name", "- @class any" },
},
},
}

View File

@@ -2,95 +2,99 @@ local ts_utils = require("nvim-treesitter.ts_utils")
return {
-- Search for these nodes
parent = { "function_definition", "class_definition" },
parent = { func = { "function_definition" }, class = { "class_definition" } },
-- Traverse down these nodes and extract the information as necessary
data = {
["function_definition"] = {
["0"] = {
extract = function(node)
local results = {}
func = {
["function_definition"] = {
["0"] = {
extract = function(node)
local results = {}
local tree = {
{
retrieve = "all",
node_type = "parameters",
subtree = {
{ retrieve = "all", node_type = "identifier", extract = true },
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 = "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 },
{
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 nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local res = neogen.utilities.extractors:extract_from_matched(nodes)
results.parameters = res.identifier
results.return_statement = res.return_statement
return results
end,
results.parameters = res.identifier
results.return_statement = res.return_statement
return results
end,
},
},
},
["class_definition"] = {
["2"] = {
match = "block",
class = {
["class_definition"] = {
["2"] = {
match = "block",
extract = function(node)
local results = {}
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 },
extract = function(node)
local results = {}
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 nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
results.attributes = {}
for _, assignment in pairs(nodes["assignment"]) do
local left_side = assignment:field("left")[1]
local left_attribute = left_side:field("attribute")[1]
table.insert(results.attributes, ts_utils.get_node_text(left_attribute)[1])
end
results.attributes = {}
for _, assignment in pairs(nodes["assignment"]) do
local left_side = assignment:field("left")[1]
local left_attribute = left_side:field("attribute")[1]
table.insert(results.attributes, ts_utils.get_node_text(left_attribute)[1])
end
return results
end,
return results
end,
},
},
},
},