First commit

lua/neogen.lua

@@ -0,0 +1,60 @@
+local ts_utils = require("nvim-treesitter.ts_utils")
+local ts_query  = require("nvim-treesitter.query")
+
+local M = {}
+
+M.generate = function (t, t2)
+    local comment = {}
+    local query = [[ 
+    (function (parameters) @params)
+    (function (return_statement) @return)
+    (function_definition (parameters) @params)
+    ]]
+    
+    -- Try to find the upper function
+    local cursor = ts_utils.get_node_at_cursor(0)
+    local function_node = cursor
+    while function_node ~= nil do
+        if function_node:type() == "function_definition" then break end
+        if function_node:type() == "function" then break end
+        function_node = function_node:parent()
+    end
+    local line = ts_utils.get_node_range(function_node)
+
+    local returned = vim.treesitter.parse_query("lua", query)
+    for id, node in returned:iter_captures(function_node) do -- For each found query
+
+        -- Try to add params
+        if returned.captures[id] == "params" then 
+            local params = ts_utils.get_node_text(node)[1]:sub(2,-2)
+            for p in string.gmatch(params, '[^,]+') do
+                p = p:gsub("%s+", "") -- remove trailing spaces
+                table.insert(comment, "---@param " .. p .. " ")
+            end
+        end
+
+        -- Try to add return statement
+        if returned.captures[id] == "return" then
+            table.insert(comment, "---@return ")
+        end
+    end
+    
+    -- Write on top of function
+    vim.fn.append(line, comment)
+    vim.fn.cursor(line+1, #comment[1])
+    vim.api.nvim_command('startinsert!')
+
+end
+
+function M.generate_command()
+  vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").generate()')
+end
+
+M.setup = function(opts)
+    local config = opts or {}
+    if config.enabled == true then
+       M.generate_command()
+    end
+end
+
+return M