68 lines
2.1 KiB
Lua
68 lines
2.1 KiB
Lua
local ts_utils = require("nvim-treesitter.ts_utils")
|
|
|
|
neogen = {}
|
|
|
|
local configuration = require('neogen.config')
|
|
|
|
neogen.generate = function ()
|
|
local comment = {}
|
|
|
|
-- 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
|
|
if function_node:type() == "local_function" then break end
|
|
function_node = function_node:parent()
|
|
end
|
|
local line = ts_utils.get_node_range(function_node)
|
|
|
|
-- find the starting position in the line function
|
|
local line_content = vim.api.nvim_buf_get_lines(0, line, line+1, false)[1]
|
|
local offset = line_content:match("^%s+") or ""
|
|
|
|
local return_comment = offset .. "---@return "
|
|
local param_comment = offset .. "---@param "
|
|
|
|
-- Parse and iterate over each found query
|
|
local returned = vim.treesitter.get_query("lua", "neogen")
|
|
for id, node in returned:iter_captures(function_node) do
|
|
|
|
-- 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_comment .. p .. " ")
|
|
end
|
|
end
|
|
|
|
-- Try to add return statement
|
|
if returned.captures[id] == "return" then
|
|
table.insert(comment, return_comment)
|
|
end
|
|
end
|
|
|
|
-- At the end, add description annotation
|
|
table.insert(comment, 1, offset .. "---")
|
|
|
|
if #comment == 0 then return 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 neogen.generate_command()
|
|
vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").generate()')
|
|
end
|
|
|
|
neogen.setup = function(opts)
|
|
local config = opts or configuration
|
|
if config.enabled == true then neogen.generate_command() end
|
|
end
|
|
|
|
return neogen
|