Add return statement for python function

lua/neogen/configurations/python.lua

@@ -7,16 +7,28 @@ return {
     -- Traverse down these nodes and extract the information as necessary
     data = {
         ["function_definition"] = {
-            ["2"] = {
+            ["0"] = {
-                match = "parameters",
+                extract = function (node)
+                    local results = {
 
-                extract = function(node)
+                    }
-                    local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(node)
+                    local regular_params = neogen.utilities.extractors:extract_children_from({ ["2"] = "extract" }, "identifier")(node)
+
+                    local body = neogen.utilities.nodes:matching_child_nodes(node, "block")[1]
+                    if body == nil then
+                        return
+                    end
+
+                    local return_statement = neogen.utilities.nodes:matching_child_nodes(body, "return_statement")
+                    if #return_statement == 0 then
+                        return_statement = nil
+                    end
 
                     return {
                         parameters = regular_params,
+                        return_statement = return_statement
                     }
-                end,
+                end
             },
         },
         ["class_definition"] = {
@@ -69,6 +81,7 @@ return {
             { nil, '"""' },
             { "parameters", "\t%s: ", { before_first_item = "Args: " } },
             { "attributes", "\t%s: ", { before_first_item = "Attributes: " } },
+            { "return_statement", "", { before_first_item = "Returns: " } },
             { nil, '"""' },
         },
     },

lua/neogen/granulators/default.lua

@@ -13,7 +13,13 @@ neogen.default_granulator = function(parent_node, node_data)
         if vim.tbl_contains(matches, parent_node:type()) then
             -- For each child_data in the matched parent node
             for i, data in pairs(child_data) do
-                local child_node = parent_node:named_child(tonumber(i) - 1)
+                local child_node
+
+                if tonumber(i) == 0 then
+                   child_node = parent_node
+                else
+                    child_node = parent_node:named_child(tonumber(i) - 1)
+                end
 
                 if not child_node then
                     return