fix: python treesitter class definition query

fynnfluegge · Feb 7, 2024 · 55cef08 · 55cef08
1 parent fa859c1
commit 55cef08
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/doc_comments_ai/treesitter/treesitter_py.py b/doc_comments_ai/treesitter/treesitter_py.py
@@ -34,6 +34,11 @@ def _query_all_methods(self, node: tree_sitter.Node):
         for child in node.children:
             if child.type == self.method_declaration_identifier:
                 methods.append(child)
+            if child.type == "class_definition":
+                class_body = child.children[-1]
+                for child_node in class_body.children:
+                    if child_node.type == self.method_declaration_identifier:
+                        methods.append(child_node)
         return methods
 
     def _query_doc_comment(self, node: tree_sitter.Node):