feat(py) Better copy of nodes to generator (#106)

This commit is contained in:
danymat
2022-10-21 16:29:44 +02:00
parent 778dfff282
commit 0f70a414fa
4 changed files with 79 additions and 22 deletions

View File

@@ -22,18 +22,23 @@ return {
["function_definition"] = { ["function_definition"] = {
["0"] = { ["0"] = {
extract = function(node) extract = function(node)
local results = {}
local tree = { local tree = {
{ {
retrieve = "all", retrieve = "all",
node_type = "parameters", node_type = "parameters",
subtree = { subtree = {
{ retrieve = "all", node_type = "identifier", extract = true }, { retrieve = "all", node_type = "identifier", extract = true, as = i.Parameter },
{ {
retrieve = "all", retrieve = "all",
node_type = "default_parameter", node_type = "default_parameter",
subtree = { { retrieve = "all", node_type = "identifier", extract = true } }, subtree = {
{
retrieve = "all",
node_type = "identifier",
extract = true,
as = i.Parameter,
},
},
}, },
{ {
retrieve = "all", retrieve = "all",
@@ -45,7 +50,14 @@ return {
retrieve = "all", retrieve = "all",
node_type = "typed_default_parameter", node_type = "typed_default_parameter",
extract = true, extract = true,
subtree = { { retrieve = "all", node_type = "identifier", extract = true } }, subtree = {
{
retrieve = "all",
node_type = "identifier",
extract = true,
as = i.Tparam,
},
},
}, },
{ {
retrieve = "first", retrieve = "first",
@@ -70,6 +82,7 @@ return {
node_type = "return_statement", node_type = "return_statement",
recursive = true, recursive = true,
extract = true, extract = true,
as = i.Return,
}, },
{ {
retrieve = "all", retrieve = "all",
@@ -94,8 +107,9 @@ return {
}, },
} }
local nodes = nodes_utils:matching_nodes_from(node, tree) local nodes = nodes_utils:matching_nodes_from(node, tree)
local temp = {}
if nodes[i.Tparam] then if nodes[i.Tparam] then
results[i.Tparam] = {} temp[i.Tparam] = {}
for _, n in pairs(nodes[i.Tparam]) do for _, n in pairs(nodes[i.Tparam]) do
local type_subtree = { local type_subtree = {
{ retrieve = "all", node_type = "identifier", extract = true, as = i.Parameter }, { retrieve = "all", node_type = "identifier", extract = true, as = i.Parameter },
@@ -103,14 +117,15 @@ return {
} }
local typed_parameters = nodes_utils:matching_nodes_from(n, type_subtree) local typed_parameters = nodes_utils:matching_nodes_from(n, type_subtree)
typed_parameters = extractors:extract_from_matched(typed_parameters) typed_parameters = extractors:extract_from_matched(typed_parameters)
table.insert(results[i.Tparam], typed_parameters) table.insert(temp[i.Tparam], typed_parameters)
end end
end end
local res = extractors:extract_from_matched(nodes) local res = extractors:extract_from_matched(nodes)
res[i.Tparam] = temp[i.Tparam]
-- Return type hints takes precedence over all other types for generating template -- Return type hints takes precedence over all other types for generating template
if res[i.ReturnTypeHint] then if res[i.ReturnTypeHint] then
res["return_statement"] = nil res[i.HasReturn] = nil
if res[i.ReturnTypeHint][1] == "None" then if res[i.ReturnTypeHint][1] == "None" then
res[i.ReturnTypeHint] = nil res[i.ReturnTypeHint] = nil
end end
@@ -136,17 +151,31 @@ return {
end end
end end
results[i.HasParameter] = (res.typed_parameter or res.identifier) and { true } or nil local results = helpers.copy({
results[i.Type] = res.type [i.HasParameter] = function(t)
results[i.Parameter] = res.identifier return t[i.Parameter] and { true } or nil
results[i.Return] = res.return_statement end,
results[i.ReturnTypeHint] = res[i.ReturnTypeHint] [i.Type] = true,
results[i.HasReturn] = (res.return_statement or res.anonymous_return or res[i.ReturnTypeHint]) [i.Parameter] = function(t)
and { true } return t[i.Parameter]
or nil end,
results[i.ArbitraryArgs] = res[i.ArbitraryArgs] [i.Return] = true,
results[i.Kwargs] = res[i.Kwargs] [i.HasReturn] = true,
results[i.Throw] = res[i.Throw] [i.ReturnTypeHint] = true,
[i.ArbitraryArgs] = true,
[i.Kwargs] = true,
[i.Throw] = true,
[i.Tparam] = true,
}, res) or {}
-- Generates a "flag" return
results[i.HasReturn] = (results[i.ReturnTypeHint] or results[i.Return]) and { true } or nil
-- Removes generation for returns that are not typed
if results[i.ReturnTypeHint] then
results[i.Return] = nil
end
return results return results
end, end,
}, },

View File

@@ -27,7 +27,7 @@ return {
}, },
{ i.ClassAttribute, ":param %s: $1" }, { i.ClassAttribute, ":param %s: $1" },
{ i.Throw, ":raises %s: $1", { type = { "func" } } }, { i.Throw, ":raises %s: $1", { type = { "func" } } },
{ i.HasReturn, ":return: $1", { type = { "func" } } }, { i.Return, ":return: $1", { type = { "func" }, after_each = ":rtype: $1" } },
{ i.HasReturn, ":rtype: $1", { type = { "func" } } }, { i.ReturnTypeHint, ":return: $1", { type = { "func" } } },
{ nil, '"""' }, { nil, '"""' },
} }

View File

@@ -32,4 +32,32 @@ return {
get_node_text = function(node, bufnr) get_node_text = function(node, bufnr)
return vim.split(vim.treesitter.query.get_node_text(node, bufnr or 0), "\n") return vim.split(vim.treesitter.query.get_node_text(node, bufnr or 0), "\n")
end, end,
--- Copies a table to another table depending of the parameters that we want to expose
---TODO: create a doc for the table structure
---@param rules table the rules that we want to execute
---@param table table the table to copy
---@return table?
---@private
copy = function(rules, table)
P(rules, table)
local copy = {}
for parameter, rule in pairs(rules) do
local parameter_value = table[parameter]
if parameter_value then
if type(rule) == "function" then
copy[parameter] = vim.tbl_deep_extend("error", rule(table), copy[parameter] or {})
elseif rule == true and parameter_value ~= nil then
copy[parameter] = parameter_value
else
vim.notify("Incorrect rule format for parameter " .. parameter, vim.log.levels.ERROR)
return
end
end
end
return copy
end,
} }

View File

@@ -57,7 +57,7 @@ return {
--- @param tree table a nested table : { retrieve = "all|first", node_type = node_name, subtree = tree, recursive = true } --- @param tree table a nested table : { retrieve = "all|first", node_type = node_name, subtree = tree, recursive = true }
--- If you want to extract the node, do not specify the subtree and instead: extract = true --- If you want to extract the node, do not specify the subtree and instead: extract = true
--- Optional: you can specify position = number instead of retrieve, and it will fetch the child node at position number --- Optional: you can specify position = number instead of retrieve, and it will fetch the child node at position number
--- @param result table the table of results --- @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 --- @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) matching_nodes_from = function(self, parent, tree, result)
result = result or {} result = result or {}