From 55cef084c504dd38067850f82d5a28804e0af6cc Mon Sep 17 00:00:00 2001 From: fynnfluegge Date: Wed, 7 Feb 2024 22:51:50 +0100 Subject: [PATCH] fix: python treesitter class definition query --- doc_comments_ai/treesitter/treesitter_py.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/doc_comments_ai/treesitter/treesitter_py.py b/doc_comments_ai/treesitter/treesitter_py.py index 84b0637..eee3c1e 100644 --- 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):