grep_ast_pattern_search_with_tree_sitter_context.py

1from grep_ast import TreeContext, filename_to_lang
2from tree_sitter_languages import get_parser
3
4def quickstart():
5    file_path = "example.py"
6    query_text = "def " # Search for function definitions
7
8    # Read the file content
9    with open(file_path, "r") as f:
10        code = f.read()
11
12    # Initialize the parser for the specific language
13    lang = filename_to_lang(file_path)
14    parser = get_parser(lang)
15    tree = parser.parse(bytes(code, "utf8"))
16
17    # Create a TreeContext object
18    tc = TreeContext(file_path, code, tree.root_node)
19
20    # Find lines matching the query
21    for i, line in enumerate(code.splitlines()):
22        if query_text in line:
23            tc.grep_lines.add(i)
24
25    # Expand the context to include functional structures (classes, functions, etc.)
26    tc.add_context()
27
28    # Output the result with syntax highlighting/context
29    print(tc.format())
30
31if __name__ == "__main__":
32    quickstart()