Merge pull request #116 from ssiyad/autodetect_type

Calling `:Neogen` will try to find the best type used to generate annotations (#116)

It'll recursively go up the syntax tree from the cursor position. 
For example, if a function is defined inside class and the cursor is inside the function,
the annotation will be generated for the function.

README.md

@@ -69,10 +69,10 @@ use {
 - If you want to keep it simple, you can use the `:Neogen` command:
 
 ```vim
-" will generate annotation for the function you're inside
+" will generate annotation for the function, class or other relevant type you're currently in
 :Neogen
-" or you can force a certain type of annotation.
+" or you can force a certain type of annotation with `:Neogen <TYPE>`
-" It'll find the next upper node that matches the type
+" It'll find the next upper node that matches the type `TYPE`
 " E.g if you're on a method of a class and do `:Neogen class`, it'll find the class declaration and generate the annotation.
 :Neogen func|class|type|...
 ```

doc/neogen.txt

@@ -150,8 +150,8 @@ Neogen will go until the start of the function and start annotating for you.
 
 Parameters~
 {opts} `(table)` Optional configs to change default behaviour of generation.
- - {opts.type} `(string, default: "func")` Which type we are trying to use for generating annotations.
+ - {opts.type} `(string, default: "any")` Which type we are trying to use for generating annotations.
-   Currently supported: `func`, `class`, `type`, `file`
+   Currently supported: `any`, `func`, `class`, `type`, `file`
  - {opts.annotation_convention} `(table)` convention to use for generating annotations.
    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.

lua/neogen/generator.lua

@@ -16,6 +16,8 @@ local nodes = require("neogen.utilities.nodes")
 local default_locator = require("neogen.locators.default")
 local snippet = require("neogen.snippet")
 local JUMP_TEXT = "$1"
+local ANY_TYPE = "any"
+local all_types_map, reverse_type_map, populated_filetypes = {}, {}, {}
 
 local function todo_text(type)
     local i = require("neogen.types.template").item
@@ -38,15 +40,62 @@ local function todo_text(type)
     })[type] or todo["description"]
 end
 
-local function get_parent_node(filetype, typ, language)
+--- Create per filetype table of node types
+---@param filetype string
+---@param language any
+local function populate_filetype(filetype, language)
+    for lang_type, set in pairs(language.parent) do
+        if all_types_map[filetype] == nil then
+            all_types_map[filetype] = {}
+        end
+
+        if reverse_type_map[filetype] == nil then
+            reverse_type_map[filetype] = {}
+        end
+
+        for _, v in ipairs(set) do
+            table.insert(all_types_map[filetype], v)
+            reverse_type_map[filetype][v] = lang_type
+        end
+    end
+end
+
+-- Get nearest parent node
+local function get_parent_node(filetype, node_type, language)
     local parser_name = ts_parsers.ft_to_lang(filetype)
     local parser = vim.treesitter.get_parser(0, parser_name)
     local tstree = parser:parse()[1]
     local tree = tstree:root()
+    local current_node = ts_utils.get_node_at_cursor(0)
+    local match_any = node_type == ANY_TYPE
+    local target_node, target_type
+    local locator = language.locator or default_locator
 
-    language.locator = language.locator or default_locator
+    -- Get a node. `type` could be a string or table
-    -- Use the language locator to locate one of the required parent nodes above the cursor
+    local function get_node(type)
-    return language.locator({ root = tree, current = ts_utils.get_node_at_cursor(0) }, language.parent[typ])
+        return locator({
+            root = tree,
+            current = current_node
+        }, type)
+    end
+
+    if match_any then
+        -- Populate node types for filetype if not already
+        if populated_filetypes[filetype] == nil then
+            populate_filetype(filetype, language)
+            populated_filetypes[filetype] = true
+        end
+
+        target_node = get_node(all_types_map[filetype])
+
+        -- Lookup node type from a map since type names slightly differ
+        target_type = target_node and reverse_type_map[filetype][target_node:type()]
+    else
+        target_node = get_node(language.parent[node_type])
+        target_type = node_type
+    end
+
+    return target_node, target_type
 end
 
 --- Generates the prefix according to `template` options.
@@ -190,7 +239,7 @@ local function generate_content(parent, data, template, required_type, annotatio
 end
 
 return setmetatable({}, {
-    __call = function(_, filetype, typ, return_snippet, annotation_convention)
+    __call = function(_, filetype, node_type, return_snippet, annotation_convention)
         if filetype == "" then
             notify("No filetype detected", vim.log.levels.WARN)
             return
@@ -200,12 +249,14 @@ return setmetatable({}, {
             notify("Language " .. filetype .. " not supported.", vim.log.levels.WARN)
             return
         end
-        typ = (type(typ) ~= "string" or typ == "") and "func" or typ -- Default type
+        node_type = (type(node_type) ~= "string" or node_type == "") and ANY_TYPE or node_type -- Default type
+
         local template = language.template
-        if not language.parent[typ] or not language.data[typ] or not template or not template.annotation_convention then
+        if not template or not template.annotation_convention then
-            notify("Type `" .. typ .. "` not supported", vim.log.levels.WARN)
+            notify("Language " .. filetype .. " does not support any annotation convention", vim.log.levels.WARN)
             return
         end
+
         annotation_convention = annotation_convention or {}
         if annotation_convention[filetype] and not template[annotation_convention[filetype]] then
             notify(
@@ -218,19 +269,26 @@ return setmetatable({}, {
             return
         end
 
-        local parent_node = get_parent_node(filetype, typ, language)
+        if node_type ~= ANY_TYPE then
+            if not language.parent[node_type] or not language.data[node_type] then
+                notify("Type `" .. node_type .. "` not supported", vim.log.levels.WARN)
+                return
+            end
+        end
+
+        local parent_node, node_type = get_parent_node(filetype, node_type, language)
         if not parent_node then
             return
         end
 
-        local data = granulator(parent_node, language.data[typ])
+        local data = granulator(parent_node, language.data[node_type])
 
         -- 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,
-            typ,
+            node_type,
             annotation_convention[filetype]
         )
 

lua/neogen/init.lua

@@ -148,8 +148,8 @@ neogen.configuration = {
 --- Neogen will go until the start of the function and start annotating for you.
 ---
 ---@param opts table Optional configs to change default behaviour of generation.
----  - {opts.type} `(string, default: "func")` Which type we are trying to use for generating annotations.
+---  - {opts.type} `(string, default: "any")` Which type we are trying to use for generating annotations.
----    Currently supported: `func`, `class`, `type`, `file`
+---    Currently supported: `any`, `func`, `class`, `type`, `file`
 ---  - {opts.annotation_convention} `(table)` convention to use for generating annotations.
 ---    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.
@@ -241,6 +241,11 @@ end
 ---
 --- Note: We will only document `major` and `minor` versions, not `patch` ones. (only X and Y in X.Y.z)
 ---
+--- ## 2.11.0~
+---   - Calling `:Neogen` will try to find the best type used to generate annotations (#116)
+---     It'll recursively go up the syntax tree from the cursor position. 
+---     For example, if a function is defined inside class and the cursor is inside the function,
+---     the annotation will be generated for the function.
 --- ## 2.10.0~
 ---   - Add support for Throw statements in python
 ---     Note: only active for reST template as of right now (please open an issue request for more templates)
@@ -285,7 +290,7 @@ end
 ---     with multiple annotation conventions.
 ---@tag neogen-changelog
 ---@toc_entry Changes in neogen plugin
-neogen.version = "2.10.4"
+neogen.version = "2.11.0"
 --minidoc_afterlines_end
 
 return neogen