Tree-sitter的5种强大用法

Tree-sitter是一个很棒的解析器，可以帮助我们像专业人士一样理解源代码。无论我们是想构建分析代码的工具、创建代码检查器，还是检测AI生成的代码，Tree-sitter都能提供帮助。

在本文中，我们将探讨五种实用方法来使用Tree-sitter，并提供您可以自己尝试的实践示例。让我们开始吧！

1、什么是Tree-sitter？

Tree-sitter是一个解析库，它构建源代码的结构化树。这意味着我们可以轻松地检查代码结构，而无需使用脆弱的正则表达式。它还支持多种语言，如Python、JavaScript、C、Go等！

开始之前：

确保我们已安装tree-sitter：

pip install tree-sitter==0.20.1

有关详细的设置说明，请按照此GitHub Gist操作！我们将在接下来的示例中使用以下示例代码来检查tree-sitter的用法

# example.py

# TODO: Optimize this
def add(a, b):
    return a + b

def subtract(a, b):
    """Subtract two numbers"""
    return a - b

def multiply(a, b):
    if a > b:
        return a * b
    else:
        return a * b

def main():
    print(add(2, 3))
    print(subtract(5, 2))
    print(multiply(2, 2))

2、提取函数及其名称

💭 问题： 想象一下，我们正在生成文档或分析大型代码库。如何自动列出所有顶级函数？

我们将把文件解析为语法树并过滤function_definition节点——无需正则表达式！

from tree_sitter import Language, Parser

LANG = Language('build/my-languages.so', 'python')
parser = Parser()
parser.set_language(LANG)

code = open('example.py', 'rb').read()
tree = parser.parse(code)

def find_functions(node):
    results = []
    if node.type == "function_definition":
        results.append(node)
    for child in node.children:
        results.extend(find_functions(child))
    return results

functions = find_functions(tree.root_node)
for func in functions:
    name_node = func.child_by_field_name('name')
    func_name = code[name_node.start_byte:name_node.end_byte].decode()
    print("Found function:", func_name)

输出：

Found function: add
Found function: subtract
Found function: multiply
Found function: main

2、在注释中查找TODO

💭 问题： 如果我们想跟踪团队中代码库的所有TODO怎么办？搜索纯文本是有风险的，如果我们错过注释掉的代码怎么办？

我们可以只查看comment节点，这样您永远不会错过或误识别它们。

def find_comments(node):
    todos = []
    if node.type == "comment":
        text = code[node.start_byte:node.end_byte].decode()
        if "TODO" in text:
            todos.append(text)
    for child in node.children:
        todos.extend(find_comments(child))
    return todos

todos = find_comments(tree.root_node)
print("\nFound TODOs:", todos)

输出：

Found TODOs: ['# TODO: Optimize this']

4、生成调用图

💭 问题： 假设我们需要列出每个调用的函数以确定依赖关系或死代码。

我们将遍历call节点以发现所有调用的函数。

def find_calls(node):
    calls = []
    if node.type == "call":
        func_node = node.child_by_field_name('function')
        func_name = code[func_node.start_byte:func_node.end_byte].decode()
        calls.append(func_name)
    for child in node.children:
        calls.extend(find_calls(child))
    return calls

call_list = find_calls(tree.root_node)
print("\nFunction calls:", call_list)

输出：

Function calls: ['print', 'add', 'print', 'subtract', 'print', 'multiply']

5、计算嵌套块

💭 问题： 复杂的代码更难阅读和维护。我们如何快速检查哪些函数包含嵌套的if或循环？

我们将遍历树并计算if_statement节点的深度。

def max_depth(node, depth=0):
    max_d = depth
    for child in node.children:
        max_d = max(max_d, max_depth(child, depth + (1 if child.type == "if_statement" else 0)))
    return max_d

print("\nMaximum if-statement nesting depth:", max_depth(tree.root_node))

输出：

Maximum if-statement nesting depth: 1

6、强制执行代码标准（自定义代码检查器）

💭 问题： 如果我们想强制每个函数都有适当的文档字符串怎么办？

我们可以检查每个函数的主体并验证第一条语句是否为string。

def has_docstring(func_node):
    body = func_node.child_by_field_name('body')
    first_stmt = body.children[0] if body.children else None
    return first_stmt and first_stmt.type == "expression_statement"

for func in functions:
    name_node = func.child_by_field_name('name')
    func_name = code[name_node.start_byte:name_node.end_byte].decode()
    print(f"{func_name}: has_docstring={has_docstring(func)}")

输出：

add: has_docstring=False
subtract: has_docstring=True
multiply: has_docstring=False
main: has_docstring=False

Tree-sitter不仅仅是一个解析器；它是一个用于解锁代码结构的工具包。凭借其快速、准确的语法树，我们可以：

提取有意义的信息，如函数和调用
改进质量检查，如文档字符串强制执行
衡量复杂度并轻松查找TODO
为您的团队或研究构建高级工具

这还只是皮毛；您可以将这些示例适应其他语言，或更深入地了解静态分析、重构等。

原文链接: 5 Powerful Ways to Use Tree-sitter in Your Next Project

汇智网翻译整理，转载请标明出处