From a79fa586a96f49b427f016e91770238244044ffa Mon Sep 17 00:00:00 2001 From: danymat Date: Thu, 13 Jan 2022 09:35:41 +0100 Subject: [PATCH] (py) Update configurations --- lua/neogen/configurations/python.lua | 8 +++++++- lua/neogen/utilities/nodes.lua | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/lua/neogen/configurations/python.lua b/lua/neogen/configurations/python.lua index 24ce7d1..3cd9d73 100644 --- a/lua/neogen/configurations/python.lua +++ b/lua/neogen/configurations/python.lua @@ -51,13 +51,19 @@ return { { retrieve = "all", node_type = "return_statement", extract = true }, }, }, + { + retrieve = "all", + node_type = "type", + as = "return_statement", + extract = true, + }, } local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree) local res = neogen.utilities.extractors:extract_from_matched(nodes) results.type = res.type results.parameters = res.identifier - results.return_statement = res.return_statement + results.return_statement = res.return_statement and { res.return_statement[1] } or nil return results end, }, diff --git a/lua/neogen/utilities/nodes.lua b/lua/neogen/utilities/nodes.lua index 2df3bf0..e35abd9 100644 --- a/lua/neogen/utilities/nodes.lua +++ b/lua/neogen/utilities/nodes.lua @@ -87,7 +87,7 @@ neogen.utilities.nodes = { for _, child in pairs(matched) do if subtree.extract == true then - local name = subtree.node_type or "_" + local name = subtree.as and subtree.as or (subtree.node_type or "_") if result[name] == nil then result[name] = {} end