Skip to content

Commit

Permalink
fix: python treesitter class definition query
Browse files Browse the repository at this point in the history
  • Loading branch information
fynnfluegge committed Feb 7, 2024
1 parent fa859c1 commit 55cef08
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions doc_comments_ai/treesitter/treesitter_py.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 55cef08

Please sign in to comment.