diff --git a/after/queries/lua/neogen.scm b/after/queries/lua/neogen.scm deleted file mode 100644 index 6d72a7f..0000000 --- a/after/queries/lua/neogen.scm +++ /dev/null @@ -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) diff --git a/lua/neogen.lua b/lua/neogen.lua index 01ad38d..7e67e17 100644 --- a/lua/neogen.lua +++ b/lua/neogen.lua @@ -1,67 +1,73 @@ -local ts_utils = require("nvim-treesitter.ts_utils") +local ok, ts_utils = pcall(require, "nvim-treesitter.ts_utils") +assert(ok, "neogen requires nvim-treesitter to operate :(") neogen = {} -local configuration = require('neogen.config') +require("neogen.utility") -neogen.generate = function () - local comment = {} +-- Require defaults +require("neogen.locators.default") +require("neogen.granulators.default") +require("neogen.generators.default") - -- 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 "" +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()] - local return_comment = offset .. "---@return " - local param_comment = offset .. "---@param " + if language then + language.locator = language.locator or neogen.default_locator + language.granulator = language.granulator or neogen.default_granulator + language.generator = language.generator or neogen.default_generator - -- 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 + -- Use the language locator to locate one of the required parent nodes above the cursor + local located_parent_node = language.locator({ + root = tree:root(), + current = ts_utils.get_node_at_cursor(0), + }, language.parent) - -- 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 .. " ") + if not located_parent_node then + return + end + + -- Use the language granulator to get the required content inside the node found with the locator + local data = language.granulator(located_parent_node, language.data) + + if data and not vim.tbl_isempty(data) then + -- Will try to generate the documentation from a template and the data found from the granulator + local to_place, start_column, content = language.generator( + located_parent_node, + data, + custom_template or language.template + ) + + -- Append the annotation in required place + vim.fn.append(to_place, content) + + -- Place cursor after annotations ans start editing + if neogen.configuration.input_after_comment == true then + vim.fn.cursor(to_place+1, start_column) + vim.api.nvim_command('startinsert!') + end 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) end function neogen.generate_command() - vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").generate()') + vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").auto_generate()') end neogen.setup = function(opts) - local config = opts or configuration - if config.enabled == true then neogen.generate_command() end + neogen.configuration = vim.tbl_deep_extend("keep", opts or {}, { + input_after_comment = true, + -- DEFAULT CONFIGURATION + languages = { + lua = require("neogen.configurations.lua"), + }, + }) + + neogen.generate_command() end return neogen diff --git a/lua/neogen/config.lua b/lua/neogen/config.lua deleted file mode 100644 index 78ea978..0000000 --- a/lua/neogen/config.lua +++ /dev/null @@ -1,5 +0,0 @@ -neogen.configuration = { - enabled = false, -} - -return neogen diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua new file mode 100644 index 0000000..f391678 --- /dev/null +++ b/lua/neogen/configurations/lua.lua @@ -0,0 +1,64 @@ +local ts_utils = require("nvim-treesitter.ts_utils") + +return { + -- Search for these nodes + parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" }, + + -- Traverse down these nodes and extract the information as necessary + data = { + ["function|local_function"] = { + -- Get second child from the parent node + ["2"] = { + -- It has to be of type "parameters" + match = "parameters", + + extract = function(node) + local regular_params = neogen.utility:extract_children("identifier")(node) + local varargs = neogen.utility:extract_children("spread")(node) + + return { + parameters = regular_params, + vararg = varargs, + } + end, + }, + }, + ["local_variable_declaration|field|variable_declaration"] = { + ["2"] = { + match = "function_definition", + + extract = function(node) + local regular_params = neogen.utility:extract_children_from({ + [1] = "extract", + }, "identifier")(node) + + local varargs = neogen.utility:extract_children_from({ + [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, + }, + }, + }, + + -- Custom lua locator that escapes from comments + locator = require("neogen.locators.lua"), + + -- Use default granulator and generator + granulator = nil, + generator = nil, + + template = { + { nil, "- " }, + { "parameters", "- @param %s any" }, + { "vararg", "- @vararg any" }, + { "return_statement", "- @return any" } + }, +} diff --git a/lua/neogen/generators/default.lua b/lua/neogen/generators/default.lua new file mode 100644 index 0000000..aa5e80f --- /dev/null +++ b/lua/neogen/generators/default.lua @@ -0,0 +1,54 @@ +local ts_utils = require("nvim-treesitter.ts_utils") + +---Default Generator: +---Uses the provided template to format the annotations with data found by the granulator +--- @param parent userdata the node used to generate the annotations +--- @param data table the data from the granulator, which is a set of [type] = results +--- @param template table a template from the configuration +--- @return table { line, content }, with line being the line to append the content +neogen.default_generator = function(parent, data, template) + local start_row, start_column, _, _ = ts_utils.get_node_range(parent) + local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format("")) + + if not template then + -- Default template + generated_template = { + { nil, "" }, + { "name", " @Summary " }, + { "parameters", " @Param " }, + { "return", " @Return " }, + } + elseif type(template) == "function" then + -- You can also pass a function as a template + generated_template = template(parent, commentstring, data) + else + generated_template = template + end + + local function parse_generated_template() + local result = {} + local prefix = (" "):rep(start_column) .. commentstring + + for _, values in ipairs(generated_template) do + local type = values[1] + + if not type then + table.insert(result, prefix .. values[2]:format("")) + else + if data[type] then + if #vim.tbl_values(data[type]) == 1 then + table.insert(result, prefix .. values[2]:format(data[type][1])) + else + for _, value in ipairs(data[type]) do + table.insert(result, prefix .. values[2]:format(value)) + end + end + end + end + end + + return result + end + + return start_row, start_column, parse_generated_template() +end diff --git a/lua/neogen/granulators/default.lua b/lua/neogen/granulators/default.lua new file mode 100644 index 0000000..2a373aa --- /dev/null +++ b/lua/neogen/granulators/default.lua @@ -0,0 +1,55 @@ +local ts_utils = require("nvim-treesitter.ts_utils") + +--- Tries to use the configuration to find all required content nodes from the parent node +--- @param parent_node userdata the node found by the locator +--- @param node_data table the data from configurations[lang].data +neogen.default_granulator = function(parent_node, node_data) + local result = {} + + for parent_type, child_data in pairs(node_data) do + local matches = vim.split(parent_type, "|", true) + + -- Look if the parent node is one of the matches + if vim.tbl_contains(matches, parent_node:type()) then + -- For each child_data in the matched parent node + for i, data in pairs(child_data) do + local child_node = parent_node:named_child(tonumber(i) - 1) + + if not child_node then + return + end + + if child_node:type() == data.match or not data.match then + local extract = {} + + if data.extract then + -- Extract content from it { [type] = { data } } + extract = data.extract(child_node) + + -- All extracted values are created added in result, like so: [data.type] = { extract } + if data.type then + -- Extract information into a one-dimensional array + local one_dimensional_arr = {} + + for _, values in pairs(extract) do + table.insert(one_dimensional_arr, values) + end + + result[data.type] = one_dimensional_arr + else + for type, extracted_data in pairs(extract) do + result[type] = extracted_data + end + end + else + -- if not extract function, get the text from the node (required: data.type) + extract = ts_utils.get_node_text(child_node) + result[data.type] = extract + end + end + end + end + end + + return result +end diff --git a/lua/neogen/locators/default.lua b/lua/neogen/locators/default.lua new file mode 100644 index 0000000..dcb99c9 --- /dev/null +++ b/lua/neogen/locators/default.lua @@ -0,0 +1,18 @@ +--- The default locator tries to find one of the nodes to match in the current node +--- If it does not find one, will fetch the parents until he finds one +--- @param node_info table a node informations +--- @param nodes_to_match table a list of parent nodes to match +--- @return node_info.current node one of the nodes to match directly above the given node +neogen.default_locator = function(node_info, nodes_to_match) + -- If we find one of the wanted nodes in current one, return the current node + if vim.tbl_contains(nodes_to_match, node_info.current:type()) then + return node_info.current + end + + -- Else, loop to parents until we find one of the nodes to match + while node_info.current and not vim.tbl_contains(nodes_to_match, node_info.current:type()) do + node_info.current = node_info.current:parent() + end + + return node_info.current +end diff --git a/lua/neogen/locators/lua.lua b/lua/neogen/locators/lua.lua new file mode 100644 index 0000000..07b3956 --- /dev/null +++ b/lua/neogen/locators/lua.lua @@ -0,0 +1,12 @@ +local ts_utils = require("nvim-treesitter.ts_utils") + +return function(node_info, nodes_to_match) + -- We're dealing with a lua comment and we need to escape its grasp + if node_info.current:type() == "source" then + local start_row, _, _, _ = ts_utils.get_node_range(node_info.current) + vim.api.nvim_win_set_cursor(0, { start_row, 0 }) + node_info.current = ts_utils.get_node_at_cursor() + end + + return neogen.default_locator(node_info, nodes_to_match) +end diff --git a/lua/neogen/utility.lua b/lua/neogen/utility.lua new file mode 100644 index 0000000..8fa0f8a --- /dev/null +++ b/lua/neogen/utility.lua @@ -0,0 +1,50 @@ +local ts_utils = require("nvim-treesitter.ts_utils") + +neogen.utility = { + --- Return a function to extract content of required children from a node + --- @param _ any self + --- @param name string the children we want to extract (if multiple childrens, separate each one with "|") + --- @return function cb function taking a node and getting the content of each children we want from name + extract_children = function(_, name) + return function(node) + local result = {} + local split = vim.split(name, "|", true) + + for child in node:iter_children() do + if vim.tbl_contains(split, child:type()) then + table.insert(result, ts_utils.get_node_text(child)[1]) + end + end + + return result + end + end, + + --- Extract content from specified children from a tree + --- the tree parameter can be a nested { [key] = value} with key being the + --- * key: is which children we want to extract the values from (e.g first children is 1) + --- * value: "extract" or { [key] = value }. If value is "extract", it will extract the key child node + --- Example (extract the first child node from the first child node of the parent node): + --- [1] = { + --- [1] = "extract" + --- } + --- @param tree table see description + --- @param name string the children we want to extract (if multiple children, separate each one with "|") + extract_children_from = function(self, tree, name) + return function(node) + local result = {} + + for i, subtree in ipairs(tree) do + local child_node = node:named_child(i - 1) + + if subtree == "extract" then + return self:extract_children(name)(child_node) + else + return self:extract_children_from(subtree, name)(node) + end + end + + return result + end + end, +} diff --git a/stylua.toml b/stylua.toml new file mode 100644 index 0000000..bb258b9 --- /dev/null +++ b/stylua.toml @@ -0,0 +1,6 @@ +column_width = 120 +line_endings = "Unix" +indent_type = "Spaces" +indent_width = 4 +quote_style = "AutoPreferDouble" +no_call_parentheses = false