diff --git a/lua/neogen.lua b/lua/neogen.lua index 2f30332..d52341d 100644 --- a/lua/neogen.lua +++ b/lua/neogen.lua @@ -3,34 +3,32 @@ assert(ok, "neogen requires nvim-treesitter to operate :(") neogen = { } --- TODO: Move code here -neogen.generate = function(searcher, generator) end - neogen.auto_generate = function(custom_template) vim.treesitter.get_parser(0):for_each_tree(function(tree, language_tree) - local searcher = neogen.configuration.languages[language_tree:lang()] + local language = neogen.configuration.languages[language_tree:lang()] - if searcher then - searcher.locator = searcher.locator or neogen.default_locator - searcher.granulator = searcher.granulator or neogen.default_granulator - searcher.generator = searcher.generator or neogen.default_generator + if language then + language.locator = language.locator or neogen.default_locator + language.granulator = language.granulator or neogen.default_granulator + language.generator = language.generator or neogen.default_generator - local located_parent_node = searcher.locator({ + -- Use the language locator to locate one the the required parent nodes above the cursor + local located_parent_node = language.locator({ root = tree:root(), current = ts_utils.get_node_at_cursor(0), - }, searcher.parent) + }, language.parent) if not located_parent_node then return end - local data = searcher.granulator(located_parent_node, searcher.data) + local data = language.granulator(located_parent_node, language.data) if data and not vim.tbl_isempty(data) then - local to_place, content = searcher.generator( + local to_place, content = language.generator( located_parent_node, data, - custom_template or searcher.template + custom_template or language.template ) vim.fn.append(to_place, content) diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua index 8ab997e..13cd137 100644 --- a/lua/neogen/configurations/lua.lua +++ b/lua/neogen/configurations/lua.lua @@ -1,3 +1,5 @@ +local ts_utils = require("nvim-treesitter.ts_utils") + return { -- Search for these nodes parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration"}, @@ -5,7 +7,9 @@ return { -- Traverse down these nodes and extract the information as necessary data = { ["function|local_function"] = { + -- Get second child from the parent node ["2"] = { + -- It has to be of type "parameters" match = "parameters", extract = function(node) @@ -24,13 +28,13 @@ return { match = "function_definition", extract = function(node) - local regular_params = neogen.utility:extract_children_from("identifier", { + local regular_params = neogen.utility:extract_children_from({ [1] = "extract", - })(node) + }, "identifier")(node) - local varargs = neogen.utility:extract_children_from("spread", { + local varargs = neogen.utility:extract_children_from({ [1] = "extract", - })(node) + }, "spread")(node) return { parameters = regular_params, @@ -42,16 +46,7 @@ return { }, -- Custom lua locator that escapes from comments - locator = function(node_info, nodes_to_match) - -- We're dealing with a lua comment and we need to escape its grasp - if node_info.current:type() == "source" then - local start_row, _, _, _ = ts_utils.get_node_range(node_info.current) - vim.api.nvim_win_set_cursor(0, { start_row, 0 }) - node_info.current = ts_utils.get_node_at_cursor() - end - - return neogen.default_locator(node_info, nodes_to_match) - end, + locator = require('neogen.locators.lua'), -- Use default granulator and generator granulator = nil, diff --git a/lua/neogen/granulators/default.lua b/lua/neogen/granulators/default.lua index fc5e1a4..81e6ec9 100644 --- a/lua/neogen/granulators/default.lua +++ b/lua/neogen/granulators/default.lua @@ -5,10 +5,11 @@ neogen.default_granulator = function(parent_node, node_data) for parent_type, child_data in pairs(node_data) do local matches = vim.split(parent_type, "|", true) - if vim.tbl_contains(matches, parent_node:type()) then - for i, _ in pairs(node_data[parent_type]) do - local data = child_data[i] + -- Look if the parent node is one of the matches + if vim.tbl_contains(matches, parent_node:type()) then + -- For each child_data in the matched parent node + for i, data in pairs(child_data) do local child_node = parent_node:named_child(tonumber(i) - 1) if not child_node then @@ -19,8 +20,10 @@ neogen.default_granulator = function(parent_node, node_data) local extract = {} if data.extract then + -- Extract content from it { [type] = { data } } extract = data.extract(child_node) + -- All extracted values are created added in result, like so: [data.type] = { extract } if data.type then -- Extract information into a one-dimensional array local one_dimensional_arr = {} @@ -36,6 +39,7 @@ neogen.default_granulator = function(parent_node, node_data) end end else + -- if not extract function, get the text from the node (required: data.type) extract = ts_utils.get_node_text(child_node) result[data.type] = extract end diff --git a/lua/neogen/locators/default.lua b/lua/neogen/locators/default.lua index b8cd697..dcb99c9 100644 --- a/lua/neogen/locators/default.lua +++ b/lua/neogen/locators/default.lua @@ -1,7 +1,8 @@ ----The default locator tries to find one of the nodes to match in the current node ---If it does not find one, will fetch the parents until he finds one ----@param node_info table ----@param nodes_to_match table +--- The default locator tries to find one of the nodes to match in the current node +--- If it does not find one, will fetch the parents until he finds one +--- @param node_info table a node informations +--- @param nodes_to_match table a list of parent nodes to match +--- @return node_info.current node one of the nodes to match directly above the given node neogen.default_locator = function(node_info, nodes_to_match) -- If we find one of the wanted nodes in current one, return the current node if vim.tbl_contains(nodes_to_match, node_info.current:type()) then diff --git a/lua/neogen/locators/lua.lua b/lua/neogen/locators/lua.lua new file mode 100644 index 0000000..07b3956 --- /dev/null +++ b/lua/neogen/locators/lua.lua @@ -0,0 +1,12 @@ +local ts_utils = require("nvim-treesitter.ts_utils") + +return function(node_info, nodes_to_match) + -- We're dealing with a lua comment and we need to escape its grasp + if node_info.current:type() == "source" then + local start_row, _, _, _ = ts_utils.get_node_range(node_info.current) + vim.api.nvim_win_set_cursor(0, { start_row, 0 }) + node_info.current = ts_utils.get_node_at_cursor() + end + + return neogen.default_locator(node_info, nodes_to_match) +end diff --git a/lua/neogen/utility.lua b/lua/neogen/utility.lua index de15db7..c495d5a 100644 --- a/lua/neogen/utility.lua +++ b/lua/neogen/utility.lua @@ -21,26 +21,26 @@ neogen.utility = { end, --- Extract content from specified children from a set of nodes - --- the nodes parameter can be a nested { [key] = value} with key being the + --- the tree parameter can be a nested { [key] = value} with key being the --- * key: is which children we want to extract the values from (e.g first children is 1) --- * value: "extract" or { [key] = value }. If value is "extract", it will extract the key child node --- Example (extract the first child node from the first child node of the parent node): --- [1] = { --- [1] = "extract" --- } + --- @param tree table see description --- @param name string the children we want to extract (if multiple children, separate each one with "|") - --- @param nodes table see description - extract_children_from = function(self, name, nodes) + extract_children_from = function(self, tree, name) return function(node) local result = {} - for i, value in ipairs(nodes) do + for i, subtree in ipairs(tree) do local child_node = node:named_child(i - 1) - if value == "extract" then + if subtree == "extract" then return self:extract_children(name)(child_node) else - return self:extract_children_from(name, value)(node) + return self:extract_children_from(subtree, name)(node) end end