diff --git a/README.md b/README.md index 44c65d0..7b256f8 100644 --- a/README.md +++ b/README.md @@ -62,7 +62,7 @@ You can provide some options for the generate, like so: ```lua require('neogen').generate({ - type = "func" -- the annotation type to generate. Currently supported: func, class + type = "func" -- the annotation type to generate. Currently supported: func, class, type }) ``` @@ -107,7 +107,7 @@ There is a list of supported languages and fields, with their annotation style | Language | Annotation conventions | Supported fields | |---|---|---| | lua | | | -| | Emmylua (`"emmylua"`) | `@param`, `@varargs`, `@return`, `@class` | +| | Emmylua (`"emmylua"`) | `@param`, `@varargs`, `@return`, `@class`, `@type` | | python | | | | | Google docstrings (`"google_docstrings"`) | `Args`, `Attributes`, `Returns` | | | Numpydoc (`"numpydoc"`)| `Arguments`, `Attributes`, `Returns`| diff --git a/lua/neogen.lua b/lua/neogen.lua index fdbab98..21f3eab 100644 --- a/lua/neogen.lua +++ b/lua/neogen.lua @@ -44,7 +44,7 @@ neogen.generate = function(opts) if 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, language.template) + local to_place, start_column, content = language.generator(located_parent_node, data, language.template, opts.type) -- Append the annotation in required place vim.fn.append(to_place, content) diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua index 719208a..31649ce 100644 --- a/lua/neogen/configurations/lua.lua +++ b/lua/neogen/configurations/lua.lua @@ -21,11 +21,27 @@ local common_function_extractor = function(node) } end +local get_identifier_from_var = function (node) + local tree = { + { + retrieve = "first", + node_type = "variable_declarator", + subtree = { + { retrieve = "first", node_type = "identifier", extract = true }, + }, + }, + } + local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree) + local res = neogen.utilities.extractors:extract_from_matched(nodes) + return res +end + return { -- Search for these nodes parent = { func = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" }, - class = { "local_variable_declaration" }, + class = { "local_variable_declaration", "variable_declaration" }, + type = { "local_variable_declaration", "variable_declaration" }, }, data = { @@ -47,20 +63,10 @@ return { }, }, class = { - ["local_variable_declaration"] = { + ["local_variable_declaration|variable_declaration"] = { ["0"] = { extract = function(node) - local tree = { - { - retrieve = "first", - node_type = "variable_declarator", - subtree = { - { retrieve = "first", node_type = "identifier", extract = true }, - }, - }, - } - local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree) - local res = neogen.utilities.extractors:extract_from_matched(nodes) + local res = get_identifier_from_var(node) return { class_name = res.identifier, } @@ -68,6 +74,18 @@ return { }, }, }, + type = { + ["local_variable_declaration|variable_declaration"] = { + ["0"] = { + extract = function(node) + local res = get_identifier_from_var(node) + return { + type = res.identifier, + } + end, + }, + }, + }, }, -- Custom lua locator that escapes from comments @@ -81,12 +99,13 @@ return { -- Which annotation convention to use annotation_convention = "emmylua", emmylua = { - { nil, "- " }, - { nil, "- ", { no_results = true } }, + { nil, "- ", { type = { "class", "func" } } }, -- add this string only on requested types + { nil, "- ", { no_results = true } }, -- Shows only when there's no results from the granulator { "parameters", "- @param %s any" }, { "vararg", "- @vararg any" }, { "return_statement", "- @return any" }, { "class_name", "- @class any" }, + { "type", "- @type any" }, }, }, } diff --git a/lua/neogen/generators/default.lua b/lua/neogen/generators/default.lua index 114e2df..d8f963f 100644 --- a/lua/neogen/generators/default.lua +++ b/lua/neogen/generators/default.lua @@ -47,8 +47,9 @@ end --- @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 +--- @param type string --- @return table { line, content, opts }, with line being the line to append the content -neogen.default_generator = function(parent, data, template) +neogen.default_generator = function(parent, data, template, required_type) local start_row, start_column, end_row, end_column = ts_utils.get_node_range(parent) local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format("")) @@ -85,28 +86,32 @@ neogen.default_generator = function(parent, data, template) for _, values in ipairs(generated_template) do local type = values[1] local formatted_string = values[2] - local opts = values[3] or {} + local opts = values[3] or { + type = { required_type } + } -- Will append the item before all their nodes if opts.before_first_item and data[type] then result = add_values_to_result(result, opts.before_first_item, prefix) end - -- If there is no data returned, will append the string with opts.no_results - if opts.no_results == true and vim.tbl_isempty(data) then - local inserted = conditional_prefix_inserter(prefix, formatted_string) - table.insert(result, inserted) - else - -- append the output as is - if type == nil and opts.no_results ~= true and not vim.tbl_isempty(data) then - local inserted = conditional_prefix_inserter(prefix, formatted_string:format("")) + if opts.type and vim.tbl_contains(opts.type, required_type) then + -- If there is no data returned, will append the string with opts.no_results + if opts.no_results == true and vim.tbl_isempty(data) then + local inserted = conditional_prefix_inserter(prefix, formatted_string) table.insert(result, inserted) else - -- Format the output with the corresponding data - if data[type] then - for _, value in ipairs(data[type]) do - local inserted = conditional_prefix_inserter(prefix, formatted_string:format(value)) - table.insert(result, inserted) + -- append the output as is + if type == nil and opts.no_results ~= true and not vim.tbl_isempty(data) then + local inserted = conditional_prefix_inserter(prefix, formatted_string:format("")) + table.insert(result, inserted) + else + -- Format the output with the corresponding data + if data[type] then + for _, value in ipairs(data[type]) do + local inserted = conditional_prefix_inserter(prefix, formatted_string:format(value)) + table.insert(result, inserted) + end end end end