diff --git a/README.md b/README.md index bd29bb21..88b1daba 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,23 @@ local opts = { noremap = true, silent = true } vim.api.nvim_set_keymap("n", "nc", ":lua require('neogen').generate({ type = 'class' })", opts) ``` +If you'd like to generate only part of a docstring, add `sections`. + +```lua +require('neogen').generate({ + type = "func" -- the annotation type to generate. Currently supported: func, class, type, file + sections = {"parameter", "return"} +}) +``` + +Example sections: +``` +parameter +return +throw +yield +``` + ### Snippet support We added snippet support, and we provide defaults for some snippet engines. diff --git a/doc/neogen.txt b/doc/neogen.txt index 654f613a..1455f2fa 100644 --- a/doc/neogen.txt +++ b/doc/neogen.txt @@ -182,12 +182,15 @@ Feel free to submit a PR, I will be happy to help you ! We use semantic versioning ! (https://semver.org) Here is the current Neogen version: >lua - neogen.version = "2.19.4" + + neogen.version = "2.20.0" < # Changelog~ Note: We will only document `major` and `minor` versions, not `patch` ones. (only X and Y in X.Y.z) +## 2.20.0~ + - Added `sections` support for docstring generation (#185) ## 2.19.0~ - Add support for julia (`julia`) ! (#185) ## 2.18.0~ @@ -464,4 +467,4 @@ If not specified, will use this line for all types. {required} `(string)` If specified, is used in if the first field of the table is a `table` (example above) - vim:tw=78:ts=8:noet:ft=help:norl: \ No newline at end of file + vim:tw=78:ts=8:noet:ft=help:norl: diff --git a/lua/neogen/configurations/python.lua b/lua/neogen/configurations/python.lua index 02d2c4b6..2364cc21 100644 --- a/lua/neogen/configurations/python.lua +++ b/lua/neogen/configurations/python.lua @@ -460,8 +460,8 @@ return { }, }, }, + locator = require("neogen.locators.python"), -- Use default granulator and generator - locator = nil, granulator = nil, generator = nil, template = template diff --git a/lua/neogen/generator.lua b/lua/neogen/generator.lua index 67918f63..a22170f7 100644 --- a/lua/neogen/generator.lua +++ b/lua/neogen/generator.lua @@ -152,8 +152,13 @@ end ---@param template table a template from the configuration ---@param required_type string ---@param annotation_convention string +---@param sections string[] | string the parts of a docstring to create ---@return table { line, content }, with line being the line to append the content -local function generate_content(parent, data, template, required_type, annotation_convention) +local function generate_content(parent, data, template, required_type, annotation_convention, partial) + if partial == nil then + partial = false + end + local row, col = get_place_pos(parent, template.position, template.append, required_type) local commentstring = vim.trim(vim.bo.commentstring:format("")) @@ -188,7 +193,7 @@ local function generate_content(parent, data, template, required_type, annotatio end local ins_type = type(inserted_type) - if ins_type == "nil" then + if not partial and ins_type == "nil" then local no_data = vim.tbl_isempty(data) if opts.no_results then if no_data then @@ -231,11 +236,110 @@ local function generate_content(parent, data, template, required_type, annotatio end end + if partial then + local index = 1 + + while result[index] == "" do + table.remove(result, index) + end + end + return row, result, default_text end +--- Interpret all `sections` into Neogen-compatible section names. +--- Each section name is a partial match. e.g. "parameter" will match +--- "has_parameter" and "parameters". +---@param sections string[] | string A user's desired parts of a docstring to create. +---@return string[] # The resolved section names. +--- +local function expand_sections(sections) + local function has_match(expression, options) + for _, option in ipairs(options) do + if option:match(expression) then + return true + end + end + + return false + end + + local i = require("neogen.types.template").item + + if type(sections) == "string" then + sections = {sections} + end + + local has_parameter = false + local parameters = { + i.ArbitraryArgs, + i.HasParameter, + i.Kwargs, + i.Parameter, + i.Tparam, + i.Vararg, + } + + local has_return = false + local returns = { + i.HasReturn, + i.Return, + i.ReturnAnonym, + i.ReturnTypeHint, + } + + local has_throw = false + local throws = { i.HasThrow, i.Throw } + + local has_yield = false + local yields = { i.HasYield, i.Yield } + + local output = {} + + for _, section in ipairs(sections) do + if not has_parameter and has_match(section, parameters) then + vim.list_extend(output, parameters) + has_parameter = true + end + + if not has_return and has_match(section, returns) then + vim.list_extend(output, returns) + has_return = true + end + + if not has_throw and has_match(section, throws) then + vim.list_extend(output, throws) + has_throw = true + end + + if not has_yield and has_match(section, yields) then + vim.list_extend(output, yields) + has_yield = true + end + end + + return output +end + +--- Remove `data` that is not present in `sections`. +---@param data table the data from configurations[lang].data +---@param sections string[] Any part of `data` not found in `sections` will be omitted. +---@return table # The filtered output of `data`. +local function filter_by_sections(data, sections) + local output = {} + + for key, value in pairs(data) do + if vim.tbl_contains(sections, key) then + output[key] = value + end + end + + return output +end + + return setmetatable({}, { - __call = function(_, filetype, node_type, return_snippet, annotation_convention) + __call = function(_, filetype, node_type, return_snippet, annotation_convention, sections) if filetype == "" then notify("No filetype detected", vim.log.levels.WARN) return @@ -279,9 +383,21 @@ return setmetatable({}, { local data = granulator(parent_node, language.data[node_type]) + local partial = false + + if sections ~= nil then + partial = true + sections = expand_sections(sections) + data = filter_by_sections(data, sections) + end + -- Will try to generate the documentation from a template and the data found from the granulator local row, template_content, default_text = - generate_content(parent_node, data, template, node_type, annotation_convention[filetype]) + generate_content(parent_node, data, template, node_type, annotation_convention[filetype], partial) + + if partial then + row = vim.api.nvim_win_get_cursor(0)[1] + end local content = {} local marks_pos = {} diff --git a/lua/neogen/init.lua b/lua/neogen/init.lua index 160d6c7b..43c2f2b7 100644 --- a/lua/neogen/init.lua +++ b/lua/neogen/init.lua @@ -154,6 +154,7 @@ neogen.configuration = { --- This is language specific. For example, `generate({ annotation_convention = { python = 'numpydoc' } })` --- If no convention is specified for a specific language, it'll use the default annotation convention for the language. --- - {opts.return_snippet} `boolean` if true, will return 3 values from the function call. +--- - {opts.sections} `string | string[]` if provided, only these parts of a docstring are created. --- This option is useful if you want to get the snippet to use with a unsupported snippet engine --- Below are the returned values: --- - 1: (type: `string[]`) the resulting lsp snippet @@ -167,7 +168,7 @@ neogen.generate = function(opts) end opts = opts or {} - return require("neogen.generator")(vim.bo.filetype, opts.type, opts.return_snippet, opts.annotation_convention) + return require("neogen.generator")(vim.bo.filetype, opts.type, opts.return_snippet, opts.annotation_convention, opts.sections) end -- Expose more API ============================================================ @@ -309,7 +310,7 @@ end --- with multiple annotation conventions. ---@tag neogen-changelog ---@toc_entry Changes in neogen plugin -neogen.version = "2.19.4" +neogen.version = "2.20.0" --minidoc_afterlines_end return neogen diff --git a/lua/neogen/locators/python.lua b/lua/neogen/locators/python.lua new file mode 100644 index 00000000..c24feaca --- /dev/null +++ b/lua/neogen/locators/python.lua @@ -0,0 +1,17 @@ +---@param node_info Neogen.node_info +---@param nodes_to_match TSNode[] +---@return TSNode? +return function(node_info, nodes_to_match) + local current = node_info.current + local parents = { "class_definition", "function_definition", "module" } + + while current do + if vim.tbl_contains(parents, current:type()) then + return current + end + + current = current:parent() + end + + return nil +end diff --git a/tests/neogen/python_google_spec.lua b/tests/neogen/python_google_spec.lua index f69ec026..b85bd47f 100644 --- a/tests/neogen/python_google_spec.lua +++ b/tests/neogen/python_google_spec.lua @@ -4,8 +4,12 @@ local specs = require('tests.utils.specs') -local function make_google_docstrings(source) - return specs.make_docstring(source, 'python', { annotation_convention = { python = 'google_docstrings' } }) +--- Parse `source` and expand all docstrings with Neogen. +---@param source string Pseudo-source-code to call. It contains `"|cursor|"` which is the expected user position. +---@param sections (string[] | string)? Parts of the docstring to create. e.g. `"parameters"`. +---@return string # The real, Neogen-generated source code result. +local function make_google_docstrings(source, sections) + return specs.make_docstring(source, 'python', { annotation_convention = { python = 'google_docstrings' }, sections=sections }) end describe("python: google_docstrings", function() @@ -863,4 +867,129 @@ describe("python: google_docstrings", function() assert.equal(expected, result) end) end) + + describe("sections", function() + it("works even with no sections are given", function() + local source = [[ + def foo(thing):|cursor| + if thing: + yield 10 + yield 20 + yield 30 + else: + yield 0 + + for _ in range(10): + yield + ]] + + local expected = [[ + def foo(thing): + """[TODO:description] + + Args: + thing ([TODO:parameter]): [TODO:description] + + Yields: + [TODO:description] + """ + if thing: + yield 10 + yield 20 + yield 30 + else: + yield 0 + + for _ in range(10): + yield + ]] + + local result = make_google_docstrings(source) + + assert.equal(expected, result) + end) + + it("works with 1 section", function() + local source = [[ + def foo(thing): + """An existing docstring. + |cursor| + """ + if thing: + yield 10 + yield 20 + yield 30 + else: + yield 0 + + for _ in range(10): + yield + ]] + + local expected = [[ + def foo(thing): + """An existing docstring. + + Yields: + [TODO:description] + """ + if thing: + yield 10 + yield 20 + yield 30 + else: + yield 0 + + for _ in range(10): + yield + ]] + + local result = make_google_docstrings(source, {"yield"}) + + assert.equal(expected, result) + end) + + it("works with 2+ sections", function() + local source = [[ + def foo(thing): + """An existing docstring. + |cursor| + """ + if thing: + yield 10 + yield 20 + yield 30 + else: + yield 0 + + for _ in range(10): + yield + ]] + + local expected = [[ + def foo(thing): + """An existing docstring. + + Args: + thing ([TODO:parameter]): [TODO:description] + + Yields: + [TODO:description] + """ + if thing: + yield 10 + yield 20 + yield 30 + else: + yield 0 + + for _ in range(10): + yield + ]] + + local result = make_google_docstrings(source, {"parameters", "yield"}) + + assert.equal(expected, result) + end) + end) end)