Added support for generating sections of documentation

Added python tests

Added Python tests

Fixed rebase conflicts

Author: ColinKennedy

README.md

@@ -105,6 +105,23 @@ local opts = { noremap = true, silent = true }
 vim.api.nvim_set_keymap("n", "<Leader>nc", ":lua require('neogen').generate({ type = 'class' })<CR>", 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.

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:
+ vim:tw=78:ts=8:noet:ft=help:norl:

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

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 = {}

lua/neogen/init.lua

@@ -167,7 +167,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 +309,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

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

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)