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

@@ -1,7 +1,15 @@
 local ok, ts_utils = pcall(require, "nvim-treesitter.ts_utils")
 assert(ok, "neogen requires nvim-treesitter to operate :(")
 
-neogen = { }
+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)
@@ -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

@@ -2,7 +2,7 @@ local ts_utils = require("nvim-treesitter.ts_utils")
 
 return {
     -- Search for these nodes
-    parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration"},
+    parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" },
 
     -- Traverse down these nodes and extract the information as necessary
     data = {
@@ -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" }
     },
 }