From 03781f23d6cfc15fc49b735b7e54cd344a10aa17 Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Mon, 23 Aug 2021 15:08:41 +0200
Subject: [PATCH] Make comment in after position more resilient

---
 lua/neogen/configurations/python.lua | 6 +++---
 lua/neogen/generators/default.lua    | 7 +++----
 lua/neogen/utility.lua               | 7 ++++---
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/lua/neogen/configurations/python.lua b/lua/neogen/configurations/python.lua
index fe96bb0..ab86c78 100644
--- a/lua/neogen/configurations/python.lua
+++ b/lua/neogen/configurations/python.lua
@@ -2,7 +2,7 @@ local ts_utils = require("nvim-treesitter.ts_utils")
 
 return {
     -- Search for these nodes
-    parent = { "function_definition" },
+    parent = { "function_definition", "class_definition" },
 
     -- Traverse down these nodes and extract the information as necessary
     data = {
@@ -28,11 +28,11 @@ return {
 
     template = {
         annotation_convention = "google_docstrings", -- required: Which annotation convention to use (default_generator)
-        append = { position = "after", offset = 4 }, -- optional: where to append the text (default_generator)
+        append = { position = "after", child_number = 3 }, -- 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)
         google_docstrings = {
             { nil, '"""' },
-            { "parameters", "\t%s: ", { before_first_item = "Args: " } },
+            { "parameters", "\t%s: ", { before_first_item = "Args: " } }, -- FIXME when no parameter is set
             { nil, '"""' },
         },
     },
diff --git a/lua/neogen/generators/default.lua b/lua/neogen/generators/default.lua
index 0be48e4..991426f 100644
--- a/lua/neogen/generators/default.lua
+++ b/lua/neogen/generators/default.lua
@@ -8,17 +8,16 @@ local ts_utils = require("nvim-treesitter.ts_utils")
 --- @return table { line, content }, with line being the line to append the content
 neogen.default_generator = function(parent, data, template)
     local start_row, start_column, end_row, end_column = ts_utils.get_node_range(parent)
-    P(ts_utils.get_node_range(parent))
     local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format(""))
 
     local row_to_place = start_row
     local col_to_place = start_column
 
     local append = template.append or {}
+
     if append.position == "after" then
-        row_to_place = end_row - 1
-        -- Add the offset if there's one
-        col_to_place = start_column + (append.offset or 0)
+        local append_at_child = parent:child(append.child_number + 1)
+        row_to_place, col_to_place, _ , _ = append_at_child:range()
     end
 
     if not template or not template.annotation_convention then
diff --git a/lua/neogen/utility.lua b/lua/neogen/utility.lua
index 8fa0f8a..df52d2a 100644
--- a/lua/neogen/utility.lua
+++ b/lua/neogen/utility.lua
@@ -10,6 +10,7 @@ neogen.utility = {
             local result = {}
             local split = vim.split(name, "|", true)
 
+
             for child in node:iter_children() do
                 if vim.tbl_contains(split, child:type()) then
                     table.insert(result, ts_utils.get_node_text(child)[1])
@@ -34,13 +35,13 @@ neogen.utility = {
         return function(node)
             local result = {}
 
-            for i, subtree in ipairs(tree) do
-                local child_node = node:named_child(i - 1)
+            for i, subtree in pairs(tree) do
+                local child_node = node:named_child(tonumber(i) - 1)
 
                 if subtree == "extract" then
                     return self:extract_children(name)(child_node)
                 else
-                    return self:extract_children_from(subtree, name)(node)
+                    return self:extract_children_from(subtree, name)(child_node)
                 end
             end