4f639a9d3a
The plugin now recognizes indentation level and will write the comments
at the same indentation level as the function.
for example:
function()
-- this is junk
-- this is junk
-- The annotation will be created here
function()
end
-- this is junk
-- this is junk
end
72 lines
2.1 KiB
Lua
72 lines
2.1 KiB
Lua
local ts_utils = require("nvim-treesitter.ts_utils")
|
|
|
|
local M = {}
|
|
|
|
M.generate = function ()
|
|
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)
|
|
|
|
-- 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.parse_query("lua", query)
|
|
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 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
|