Add return statement for lua

At the moment, it only extracts return statements from direct children
of the function_definition

after/queries/lua/neogen.scm

@@ -1,7 +0,0 @@
-(function (parameters) @params)
-(function_definition (parameters) @params)
-(local_function (parameters) @params)
-
-(function (return_statement) @return)
-(function_definition (if_statement (return_statement) @return))
-(local_function (return_statement) @return)

lua/neogen.lua

@@ -3,6 +3,14 @@ assert(ok, "neogen requires nvim-treesitter to operate :(")
 
 neogen = {}
 
+require("neogen.utility")
+
+-- Require defaults
+require("neogen.locators.default")
+require("neogen.granulators.default")
+require("neogen.generators.default")
+
+
 neogen.auto_generate = function(custom_template)
     vim.treesitter.get_parser(0):for_each_tree(function(tree, language_tree)
         local language = neogen.configuration.languages[language_tree:lang()]
@@ -48,16 +56,11 @@ neogen.setup = function(opts)
     neogen.configuration = vim.tbl_deep_extend("keep", opts or {}, {
         -- DEFAULT CONFIGURATION
         languages = {
-            lua = require('neogen.configurations.lua'),
+            lua = require("neogen.configurations.lua"),
         },
     })
 
     neogen.generate_command()
 end
 
-require('neogen.utility')
-require('neogen.locators.default')
-require('neogen.granulators.default')
-require('neogen.generators.default')
-
 return neogen

lua/neogen/configurations/lua.lua

@@ -36,10 +36,12 @@ return {
                         [1] = "extract",
                     }, "spread")(node)
 
+                    local return_statement = neogen.utility:extract_children("return_statement")(node)
 
                     return {
                         parameters = regular_params,
                         vararg = varargs,
+                        return_statement = return_statement
                     }
                 end,
             },
@@ -47,7 +49,7 @@ return {
     },
 
     -- Custom lua locator that escapes from comments
-    locator = require('neogen.locators.lua'),
+    locator = require("neogen.locators.lua"),
 
     -- Use default granulator and generator
     granulator = nil,
@@ -57,5 +59,6 @@ return {
         { nil, "- " },
         { "parameters", "- @param %s any" },
         { "vararg", "- @vararg any" },
+        { "return_statement", "- @return any" }
     },
 }