Get all Python function calls with Tree Sitter queries

tree_sitter_get_function_calls.py

from athenacommon.tree_sitter import get_language
from tree_sitter import Parser

parser = Parser()
lang = get_language('python')  # get lang whatever way if you lack athenacommon
parser.set_language(lang)

code = bytes("""def foo(qux):
    h = boo()
    for j in range(10):
        qoo(h)
    return h
""", "utf-8")
print(code.decode())

tree = parser.parse(code)

query = lang.query("""((call function: (identifier) arguments: (argument_list)) @function.call)""")

print("function calls:")
function_nodes = query.captures(tree.root_node)
for node, label in function_nodes:
    print(code[node.start_byte:node.end_byte].decode())