diff --git a/README.md b/README.md index e26d8df..01b73e2 100644 --- a/README.md +++ b/README.md @@ -114,8 +114,8 @@ data = { -- Extractor function that returns a set of TSname = values with values being of type string[] extract = function(node) - local regular_params = neogen.utility:extract_children("identifier")(node) - local varargs = neogen.utility:extract_children("spread")(node) + local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(node) + local varargs = neogen.utilities.extractors:extract_children_text("spread")(node) return { parameters = regular_params, @@ -201,8 +201,8 @@ data = { -- Extractor function that returns a set of TSname = values with values being of type string[] extract = function(node) - local regular_params = neogen.utility:extract_children("identifier")(node) - local varargs = neogen.utility:extract_children("spread")(node) + local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(node) + local varargs = neogen.utilities.extractors:extract_children_text("spread")(node) return { parameters = regular_params, @@ -217,7 +217,7 @@ data = { Notes: - If you create your own granulator, you can add any kind of parameters in the `data` field from configuration file as long as the function signature is the same provided. -- Utilities are provided in order to extract content from subnodes. You can check out their documentation in `lua/utility.lua`. +- Utilities are provided. You can check out their documentation in `lua/utilities/`. ### Generators diff --git a/lua/neogen.lua b/lua/neogen.lua index 1600628..53e1d1b 100644 --- a/lua/neogen.lua +++ b/lua/neogen.lua @@ -3,7 +3,11 @@ assert(ok, "neogen requires nvim-treesitter to operate :(") neogen = {} -require("neogen.utility") + +-- Require utilities +neogen.utilities = {} +require("neogen.utilities.extractors") +require("neogen.utilities.nodes") -- Require defaults require("neogen.locators.default") diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua index dad19fe..20bd7e4 100644 --- a/lua/neogen/configurations/lua.lua +++ b/lua/neogen/configurations/lua.lua @@ -13,8 +13,8 @@ return { match = "parameters", extract = function(node) - local regular_params = neogen.utility:extract_children("identifier")(node) - local varargs = neogen.utility:extract_children("spread")(node) + local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(node) + local varargs = neogen.utilities.extractors:extract_children_text("spread")(node) return { parameters = regular_params, @@ -28,15 +28,15 @@ return { match = "function_definition", extract = function(node) - local regular_params = neogen.utility:extract_children_from({ + local regular_params = neogen.utilities.extractors:extract_children_from({ [1] = "extract", }, "identifier")(node) - local varargs = neogen.utility:extract_children_from({ + local varargs = neogen.utilities.extractors:extract_children_from({ [1] = "extract", }, "spread")(node) - local return_statement = neogen.utility:extract_children("return_statement")(node) + local return_statement = neogen.utilities.extractors:extract_children_text("return_statement")(node) return { parameters = regular_params, diff --git a/lua/neogen/configurations/python.lua b/lua/neogen/configurations/python.lua index daf47f5..1c3bd73 100644 --- a/lua/neogen/configurations/python.lua +++ b/lua/neogen/configurations/python.lua @@ -11,7 +11,7 @@ return { match = "parameters", extract = function(node) - local regular_params = neogen.utility:extract_children("identifier")(node) + local regular_params = neogen.utilities.extractors:extract_children_text("identifier")(node) return { parameters = regular_params, @@ -19,6 +19,15 @@ return { end, }, }, + ["class_definition"] = { + ["2"] = { + match = "block", + + extract = function(node) + return {} + end, + }, + }, }, -- Use default granulator and generator @@ -29,7 +38,7 @@ return { template = { annotation_convention = "google_docstrings", -- required: Which annotation convention to use (default_generator) append = { position = "after", child_name = "block" }, -- optional: where to append the text (default_generator) - use_default_comment = false, -- If you want to prefix the template with the default comment for the language (default_generator) + use_default_comment = false, -- If you want to prefix the template with the default comment for the language, e.g for python: # (default_generator) google_docstrings = { { nil, '"""' }, { "parameters", "\t%s: ", { before_first_item = "Args: " } }, diff --git a/lua/neogen/generators/default.lua b/lua/neogen/generators/default.lua index da3d80d..3e883a5 100644 --- a/lua/neogen/generators/default.lua +++ b/lua/neogen/generators/default.lua @@ -18,13 +18,10 @@ neogen.default_generator = function(parent, data, template) local append = template.append or {} if append.position == "after" then - for child in parent:iter_children() do - if child:type() == append.child_name then - row_to_place, col_to_place, _ , _ = child:range() - break - end + local child_node = neogen.utilities.nodes:first_child_node(parent, append.child_name) + if child_node ~= nil then + row_to_place, col_to_place, _ , _ = child_node:range() end - end if not template or not template.annotation_convention then diff --git a/lua/neogen/utility.lua b/lua/neogen/utilities/extractors.lua similarity index 92% rename from lua/neogen/utility.lua rename to lua/neogen/utilities/extractors.lua index df52d2a..0e64484 100644 --- a/lua/neogen/utility.lua +++ b/lua/neogen/utilities/extractors.lua @@ -1,11 +1,11 @@ local ts_utils = require("nvim-treesitter.ts_utils") -neogen.utility = { +neogen.utilities.extractors = { --- Return a function to extract content of required children from a node --- @param _ any self --- @param name string the children we want to extract (if multiple childrens, separate each one with "|") --- @return function cb function taking a node and getting the content of each children we want from name - extract_children = function(_, name) + extract_children_text = function(_, name) return function(node) local result = {} local split = vim.split(name, "|", true) @@ -39,7 +39,7 @@ neogen.utility = { local child_node = node:named_child(tonumber(i) - 1) if subtree == "extract" then - return self:extract_children(name)(child_node) + return self:extract_children_text(name)(child_node) else return self:extract_children_from(subtree, name)(child_node) end diff --git a/lua/neogen/utilities/nodes.lua b/lua/neogen/utilities/nodes.lua new file mode 100644 index 0000000..37fa1d8 --- /dev/null +++ b/lua/neogen/utilities/nodes.lua @@ -0,0 +1,14 @@ +neogen.utilities.nodes = { + --- Get first child node that match the provided node name + --- @param _ any + --- @param parent userdata the parent's node + --- @param node_name string the node type to search for + --- @return userdata node the first encountered child node + first_child_node = function (_, parent, node_name) + for child in parent:iter_children() do + if child:type() == node_name then + return child + end + end + end +}