ref(py) Better handling of class methods (#51)

lua/neogen/configurations/python.lua

@@ -1,15 +1,18 @@
 local ts_utils = require("nvim-treesitter.ts_utils")
 local nodes_utils = require("neogen.utilities.nodes")
 local extractors = require("neogen.utilities.extractors")
+local locator = require("neogen.locators.default")
 
-return {
+local parent = {
-    -- Search for these nodes
-    parent = {
     func = { "function_definition" },
     class = { "class_definition" },
     file = { "module" },
     type = { "expression_statement" },
-    },
+}
+
+return {
+    -- Search for these nodes
+    parent = parent,
 
     -- Traverse down these nodes and extract the information as necessary
     data = {
@@ -106,11 +109,10 @@ return {
                             end
                         end
 
-                        -- Remove reference to "self" if any
+                        -- Check if the function is inside a class
-                        if res.identifier then
+                        -- If so, remove reference to the first parameter (that can be `self`, `cls`, or a custom name)
-                            res.identifier = vim.tbl_filter(function(v)
+                        if res.identifier and locator({ current = node }, parent.class) then
-                                return v ~= "self"
+                            table.remove(res.identifier, 1)
-                            end, res.identifier)
                             if vim.tbl_isempty(res.identifier) then
                                 res.identifier = nil
                             end