(py) Update configurations

2022-01-13 09:35:41 +01:00
parent 7e7fbc4ff2
commit a79fa586a9
2 changed files with 8 additions and 2 deletions
--- 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,
                },