From fddd50553e1601a5e8a81a5bb6aaa55e5419ca45 Mon Sep 17 00:00:00 2001 From: Daniel Mathiot Date: Thu, 26 Aug 2021 12:28:46 +0200 Subject: [PATCH] (lua) types are now guessed for type annotations This made me to do small adjustments to the tree extractor, and it now can retrieve fixed positions from the tree. If no node_type name is specified, it will put all nodes fetched from fixed positions to ["_"] = values --- lua/neogen/configurations/javascript.lua | 10 +++--- lua/neogen/configurations/lua.lua | 34 ++++++++++++------- lua/neogen/utilities/extractors.lua | 13 ++++++-- lua/neogen/utilities/nodes.lua | 42 +++++++++++++++++------- 4 files changed, 68 insertions(+), 31 deletions(-) diff --git a/lua/neogen/configurations/javascript.lua b/lua/neogen/configurations/javascript.lua index 25ca9bb..3185b30 100644 --- a/lua/neogen/configurations/javascript.lua +++ b/lua/neogen/configurations/javascript.lua @@ -16,9 +16,9 @@ local function_tree = { } return { - parent = { - func = { "function_declaration", "expression_statement", "variable_declaration" }, - class = { "function_declaration", "expression_statement", "variable_declaration", "class_declaration" }, + parent = { + func = { "function_declaration", "expression_statement", "variable_declaration" }, + class = { "function_declaration", "expression_statement", "variable_declaration", "class_declaration" }, }, data = { @@ -64,7 +64,7 @@ return { end, }, }, - } + }, }, template = { @@ -74,7 +74,7 @@ return { jsdoc = { { nil, "/* */", { no_results = true } }, { nil, "/**" }, - { "class_tag", " * @classdesc" , { before_first_item = { " * ", " * @class" } }}, + { "class_tag", " * @classdesc", { before_first_item = { " * ", " * @class" } } }, { "parameters", " * @param {any} %s " }, { "return_statement", " * @returns {any} " }, { nil, " */" }, diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua index 1ec443e..bbd4a1b 100644 --- a/lua/neogen/configurations/lua.lua +++ b/lua/neogen/configurations/lua.lua @@ -21,7 +21,7 @@ local common_function_extractor = function(node) } end -local extract_from_var = function (node) +local extract_from_var = function(node) local tree = { { retrieve = "first", @@ -31,10 +31,13 @@ local extract_from_var = function (node) { retrieve = "first", node_type = "field_expression", extract = true }, }, }, + { + position = 2, + extract = true, + }, } local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree) - local res = neogen.utilities.extractors:extract_from_matched(nodes) - return res + return nodes end return { @@ -67,7 +70,8 @@ return { ["local_variable_declaration|variable_declaration"] = { ["0"] = { extract = function(node) - local res = extract_from_var(node) + local nodes = extract_from_var(node) + local res = neogen.utilities.extractors:extract_from_matched(nodes) return { class_name = res.identifier, } @@ -79,14 +83,20 @@ return { ["local_variable_declaration|variable_declaration"] = { ["0"] = { extract = function(node) - result = {} - local res = extract_from_var(node) + local result = {} result.type = {} - if res.identifier then - vim.list_extend(result.type, res.identifier) - end - if res.field_expression then - vim.list_extend(result.type, res.field_expression) + + local nodes = extract_from_var(node) + local res = neogen.utilities.extractors:extract_from_matched(nodes, { type = true }) + + -- We asked the extract_from_var function to find the type node at right assignment. + -- We check if it found it, or else will put `any` in the type + if res["_"] then + vim.list_extend(result.type, res["_"]) + else + if res.identifier or res.field_expression then + vim.list_extend(result.type, { "any" }) + end end return result end, @@ -112,7 +122,7 @@ return { { "vararg", "- @vararg any" }, { "return_statement", "- @return any" }, { "class_name", "- @class any" }, - { "type", "- @type any" }, + { "type", "- @type %s" }, }, }, } diff --git a/lua/neogen/utilities/extractors.lua b/lua/neogen/utilities/extractors.lua index 60a478b..6fd56a6 100644 --- a/lua/neogen/utilities/extractors.lua +++ b/lua/neogen/utilities/extractors.lua @@ -4,14 +4,23 @@ neogen.utilities.extractors = { --- Extract the content from each node from data --- @param _ any self --- @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) + extract_from_matched = function(_, data, opts) + opts = opts or {} local result = {} for k, v in pairs(data) do + local get_type = function(node) + return node:type() + end local get_text = function(node) return ts_utils.get_node_text(node)[1] end - result[k] = vim.tbl_map(get_text, v) + if opts.type == true then + result[k] = vim.tbl_map(get_type, v) + else + result[k] = vim.tbl_map(get_text, v) + end end return result end, diff --git a/lua/neogen/utilities/nodes.lua b/lua/neogen/utilities/nodes.lua index 95a3fb0..2593c9d 100644 --- a/lua/neogen/utilities/nodes.lua +++ b/lua/neogen/utilities/nodes.lua @@ -2,17 +2,26 @@ neogen.utilities.nodes = { --- Get a list of child nodes that match the provided node name --- @param _ any --- @param parent userdata the parent's node - --- @param node_name string the node type to search for (if multiple childrens, separate each one with "|") + --- @param node_name string|nil the node type to search for (if multiple childrens, separate each one with "|") --- @return table a table of nodes that matched the name matching_child_nodes = function(_, parent, node_name) local results = {} - local split = vim.split(node_name, "|", true) - - for child in parent:iter_children() do - if vim.tbl_contains(split, child:type()) then - table.insert(results, child) + -- Return all nodes if there is no node name + if node_name == nil then + for child in parent:iter_children() do + if child:named() then + table.insert(results, child) + end + end + else + local split = vim.split(node_name, "|", true) + for child in parent:iter_children() do + if vim.tbl_contains(split, child:type()) then + table.insert(results, child) + end end end + return results end, @@ -20,6 +29,7 @@ neogen.utilities.nodes = { --- @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 + --- 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 --- @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) @@ -30,20 +40,28 @@ neogen.utilities.nodes = { local matched = self:matching_child_nodes(parent, subtree.node_type) -- Only keep first matched child node + if subtree.retrieve == nil then + if type(subtree.position) == "number" then + matched = { matched[subtree.position] } + else + assert(false, "please require position if retrieve is nil") + end + end + 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] = {} + local name = subtree.node_type or "_" + if result[name] == nil then + result[name] = {} end - table.insert(result[subtree.node_type], child) + table.insert(result[name], child) else - local test = self:matching_nodes_from(child, subtree.subtree, result) - result = vim.tbl_deep_extend("keep", result, test) + local nodes = self:matching_nodes_from(child, subtree.subtree, result) + result = vim.tbl_deep_extend("keep", result, nodes) end end end