(lua) add type annotation for variables

README.md

@@ -62,7 +62,7 @@ You can provide some options for the generate, like so:
 
 ```lua
 require('neogen').generate({
-    type = "func" -- the annotation type to generate. Currently supported: func, class
+    type = "func" -- the annotation type to generate. Currently supported: func, class, type
 })
 ```
 
@@ -107,7 +107,7 @@ There is a list of supported languages and fields, with their annotation style
 | Language | Annotation conventions | Supported fields |
 |---|---|---|
 | lua | | |
-| | Emmylua (`"emmylua"`) | `@param`, `@varargs`, `@return`, `@class` |
+| | Emmylua (`"emmylua"`) | `@param`, `@varargs`, `@return`, `@class`, `@type` |
 | python | | |
 | | Google docstrings (`"google_docstrings"`) | `Args`, `Attributes`, `Returns` |
 | | Numpydoc (`"numpydoc"`)| `Arguments`, `Attributes`, `Returns`|

lua/neogen.lua

@@ -44,7 +44,7 @@ neogen.generate = function(opts)
 
             if data then
                 -- Will try to generate the documentation from a template and the data found from the granulator
-                local to_place, start_column, content = language.generator(located_parent_node, data, language.template)
+                local to_place, start_column, content = language.generator(located_parent_node, data, language.template, opts.type)
 
                 -- Append the annotation in required place
                 vim.fn.append(to_place, content)

lua/neogen/configurations/lua.lua

@@ -21,11 +21,27 @@ local common_function_extractor = function(node)
     }
 end
 
+local get_identifier_from_var = function (node)
+    local tree = {
+        {
+            retrieve = "first",
+            node_type = "variable_declarator",
+            subtree = {
+                { retrieve = "first", node_type = "identifier", extract = true },
+            },
+        },
+    }
+    local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
+    local res = neogen.utilities.extractors:extract_from_matched(nodes)
+    return res
+end
+
 return {
     -- Search for these nodes
     parent = {
         func = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" },
-        class = { "local_variable_declaration" },
+        class = { "local_variable_declaration", "variable_declaration" },
+        type = { "local_variable_declaration", "variable_declaration" },
     },
 
     data = {
@@ -47,20 +63,10 @@ return {
             },
         },
         class = {
-            ["local_variable_declaration"] = {
+            ["local_variable_declaration|variable_declaration"] = {
                 ["0"] = {
                     extract = function(node)
-                        local tree = {
+                        local res = get_identifier_from_var(node)
-                            {
-                                retrieve = "first",
-                                node_type = "variable_declarator",
-                                subtree = {
-                                    { retrieve = "first", node_type = "identifier", extract = true },
-                                },
-                            },
-                        }
-                        local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
-                        local res = neogen.utilities.extractors:extract_from_matched(nodes)
                         return {
                             class_name = res.identifier,
                         }
@@ -68,6 +74,18 @@ return {
                 },
             },
         },
+        type = {
+            ["local_variable_declaration|variable_declaration"] = {
+                ["0"] = {
+                    extract = function(node)
+                        local res = get_identifier_from_var(node)
+                        return {
+                            type = res.identifier,
+                        }
+                    end,
+                },
+            },
+        },
     },
 
     -- Custom lua locator that escapes from comments
@@ -81,12 +99,13 @@ return {
         -- Which annotation convention to use
         annotation_convention = "emmylua",
         emmylua = {
-            { nil, "- " },
+            { nil, "- ", { type = { "class", "func" } } }, -- add this string only on requested types
-            { nil, "- ", { no_results = true } },
+            { nil, "- ", { no_results = true } }, -- Shows only when there's no results from the granulator
             { "parameters", "- @param %s any" },
             { "vararg", "- @vararg any" },
             { "return_statement", "- @return any" },
             { "class_name", "- @class any" },
+            { "type", "- @type any" },
         },
     },
 }

lua/neogen/generators/default.lua

@@ -47,8 +47,9 @@ end
 --- @param parent userdata the node used to generate the annotations
 --- @param data table the data from the granulator, which is a set of [type] = results
 --- @param template table a template from the configuration
+--- @param type string
 --- @return table { line, content, opts }, with line being the line to append the content
-neogen.default_generator = function(parent, data, template)
+neogen.default_generator = function(parent, data, template, required_type)
     local start_row, start_column, end_row, end_column = ts_utils.get_node_range(parent)
     local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format(""))
 
@@ -85,28 +86,32 @@ neogen.default_generator = function(parent, data, template)
         for _, values in ipairs(generated_template) do
             local type = values[1]
             local formatted_string = values[2]
-            local opts = values[3] or {}
+            local opts = values[3] or {
+                type = { required_type }
+            }
 
             -- Will append the item before all their nodes
             if opts.before_first_item and data[type] then
                 result = add_values_to_result(result, opts.before_first_item, prefix)
             end
 
-            -- If there is no data returned, will append the string with opts.no_results
+            if opts.type and vim.tbl_contains(opts.type, required_type) then
-            if opts.no_results == true and vim.tbl_isempty(data) then
+                -- If there is no data returned, will append the string with opts.no_results
-                local inserted = conditional_prefix_inserter(prefix, formatted_string)
+                if opts.no_results == true and vim.tbl_isempty(data) then
-                table.insert(result, inserted)
+                    local inserted = conditional_prefix_inserter(prefix, formatted_string)
-            else
-                -- append the output as is
-                if type == nil and opts.no_results ~= true and not vim.tbl_isempty(data) then
-                    local inserted = conditional_prefix_inserter(prefix, formatted_string:format(""))
                     table.insert(result, inserted)
                 else
-                    -- Format the output with the corresponding data
+                    -- append the output as is
-                    if data[type] then
+                    if type == nil and opts.no_results ~= true and not vim.tbl_isempty(data) then
-                        for _, value in ipairs(data[type]) do
+                        local inserted = conditional_prefix_inserter(prefix, formatted_string:format(""))
-                            local inserted = conditional_prefix_inserter(prefix, formatted_string:format(value))
+                        table.insert(result, inserted)
-                            table.insert(result, inserted)
+                    else
+                        -- Format the output with the corresponding data
+                        if data[type] then
+                            for _, value in ipairs(data[type]) do
+                                local inserted = conditional_prefix_inserter(prefix, formatted_string:format(value))
+                                table.insert(result, inserted)
+                            end
                         end
                     end
                 end