leigh/neogen
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add annotation type autodetection

Browse Source 
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>
This commit is contained in:
Barak Michener
2022-03-02 14:49:10 -08:00
committed by Sabu Siyad
parent 91b093f11f
commit 2c454f7e68
2 changed files with 69 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
doc/neogen.txt
View File

@@ -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.
Currently supported: `func`, `class`, `type`, `file`
- {opts.type} `(string, default: "any")` Which type we are trying to use for generating annotations.
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.

77
lua/neogen/generator.lua
View File

@@ -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
-- Use the language locator to locate one of the required parent nodes above the cursor
return language.locator({ root = tree, current = ts_utils.get_node_at_cursor(0) }, language.parent[typ])
-- Get a node. `type` could be a string or table
local function get_node(type)
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
notify("Type `" .. typ .. "` not supported", vim.log.levels.WARN)
if not template or not template.annotation_convention then
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]
)