Add sections support for generating partial docstrings

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

@@ -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

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