Add annotation type autodetection

Creates a new generate() type, `any`, which contextually adds an
annotation by whichever matching type is first found in the tree around
the cursor.

So, if your cursor is in a function, it will annotate the function, if
in a class or struct definition, it'll annotate that, etc.

Co-authored-by: Sabu Siyad <hello@ssiyad.com>
Signed-off-by: Sabu Siyad <hello@ssiyad.com>

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,60 @@ local function todo_text(type)
     })[type] or todo["description"]
 end
 
-local function get_parent_node(filetype, typ, language)
+local function populate_filetype(filetype, language)
+    -- Create a table with all possible node types
+    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 +237,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,10 +247,11 @@ 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" 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("Type `" .. node_type .. "` not supported", vim.log.levels.WARN)
             return
         end
         annotation_convention = annotation_convention or {}
@@ -218,19 +266,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]
         )