fix: python treesitter class definition query

fynnfluegge · Feb 6, 2024 · deef082 · deef082
1 parent 9a3335a
commit deef082
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/codeqai/treesitter/treesitter_py.py b/codeqai/treesitter/treesitter_py.py
@@ -33,6 +33,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):