python使用tree-sitter解析cpp代码绘制出树形图可视化

1 准备环境

首先安装好两个库：

tree-sitter ：解析Cpp源代码的解析器，性能好，能快速生成 AST / CST树（具体语法树）。安装步骤具体安装参考这篇博客👉python tree-sitter 使用方法

graphviz ：可视化AST树的绘图库。通过 Graphviz官网，先下载Graphviz工具（根据自己的操作系统以及位数来选择，这里展示的是Windows 64位要下载的程序exe），然后完成本地安装。具体安装过程可以参考这篇博客👉 graphviz 可视化软件安装教程

然后在代码中引入使用

创建python环境，shell中使用conda install 或者 pip install 进行三方库的安装：

pip install tree_sitter
pip install graphviz

完成后，我们在项目当中导入两个库使用：

from tree_sitter import Language, Parser  # 解析器库
from graphviz import Digraph  # 绘图库
# ...

2 核心代码

随便用一段代码做例子，首先构造 tree-sitter 解析器，并获取到root_node节点。注意，root_node 是我们开始遍历的根节点

# 这是测试用的cpp代码片段, 斐波那契数列
cpp_code_snippet = '''
int fib(int n){
    long long f1 = 1, f2 = 1;
    int time = (n + 1) / 2;
    for (int i = 1; i < time; i++){
        f1 = f1 + f2;
        f2 = f2 + f1;
    }
    if (n % 2 == 0) return f2;
    else return f1;
}
'''

# tree-sitter CPP解析器
CPP_LANGUAGE = Language('build/my-languages.so', 'cpp')
cpp_parser = Parser()
cpp_parser.set_language(CPP_LANGUAGE)

tree = cpp_parser.parse(bytes(cpp_code_snippet, "utf8"))
root_node = tree.root_node  # 注意，root_node 才是可遍历的树节点

然后我们写一段遍历用的代码，考虑到 AST 树的层次性，这里用队列实现一个简单的广度优先遍历(BFS) ，方便后续知晓每个层次的结点数目和状态。
当然，你也可以替换成深度优先遍历(DFS)

def ast_bfs(root):
    """ traversal AST Tree stack By BFS
    """
    node_list = []  # node_list: [(node_id, node_str), ...]
    edge_list = []  # edge_list: [(parent_id, child_id), ...]

    # queue for BFS
    root_tuple = (root, None)  # 根结点 root
    queue = []
    queue.append(root_tuple)  # tuple (Node<> node, int parent_id)

    # BFS
    while queue:
      	# 出队
        (node, parent_id) = queue.pop(0)
		
		# 结点的类型
        node_type = str(node.type).strip() 
        
        # node_text 是代码内容，可以用于显示
        node_text = str(node.text.decode())
        
        # node_id 结点的ID，便与后续绘图 
        node_id = str(len(node_list))
		
		#  node_str，是你在图中希望显示的内容，这里默认“类型+代码”
        node_str = node_type + "\n" + node_text
        node_item = (node_id, node_str)  # tuple（node_id, node_string） 
        node_list.append(node_item)

        # 如果有父节点，添加一条从父节点到当前节点的边
        if parent_id is not None:
            edge_list.append((parent_id, node_id))

        # 遍历子节点入队
        for child in node.children:
            queue.append((child, node_id))

    return node_list, edge_list

最后，我们再写一段根据得到的 node_list 和 edge_list 来绘制AST树的代码

def draw_tree(graph_name: str, node_list: list[tuple], edge_list: list[tuple]):
    """ draw a tree
    """
    # 创建一个Digraph对象
    dot = Digraph(comment=graph_name)
    # 绘制结点
    for node in node_list:
        dot.node(name=node[0], label=node[1], shape="box")
    # 绘制连线
    for edge in edge_list:
        dot.edge(tail_name=edge[0], head_name=edge[1], lable="")
        
    return dot

那么得到的图dot 可以进行保存，也可以直接预览

dot.view()  # 预览结果图

最终的AST结果图如下所示：

从下图我们可以直观看到，tree-sitter 生成的是 CST（具体语法树），也就是远大买的所有细粒度符号在树中都有它的具体结点位置，包括操作符、符号、括号、分号等。

如果要得到AST（抽象语法树），最直接的一种办法就是过滤掉这些符号即可，使用正则表达式匹配也好，或者通过规律筛去也好，例如判断node.type和node.text是否完全一致，据个人观察这是绝大多数符号的特性。

3 全部代码

最后放一个整体的代码供参考学习。

from tree_sitter import Language, Parser  # 解析器库
from graphviz import Digraph  # 绘图库

# tree-sitter CPP解析器
CPP_LANGUAGE = Language('build/my-languages.so', 'cpp')
cpp_parser = Parser()
cpp_parser.set_language(CPP_LANGUAGE)

# 测试用的cpp代码片段, 斐波那契数列
cpp_code_snippet = '''
int fib(int n){
    long long f1 = 1, f2 = 1;
    int time = (n + 1) / 2;
    for (int i = 1; i < time; i++){
        f1 = f1 + f2;
        f2 = f2 + f1;
    }
    if (n % 2 == 0) return f2;
    else return f1;
}
'''

def ast_bfs(root):
    """ traversal AST Tree stack By BFS"""
    node_list = []  # node_list: [(node_id, node_type), ...]
    edge_list = []  # edge_list: [(parent_id, child_id), ...]

    # queue for BFS
    root_tuple = (root, None)  # 根节点
    queue = []
    queue.append(root_tuple)  # tuple (Node<> node, int parent_id)

    # BFS
    while queue:
        # 出队
        (node, parent_id) = queue.pop(0)

        node_type = str(node.type).strip()  # 结点的类型
        node_text = str(node.text.decode())  # 代码的内容，可以用于显示(目前没有用)
        node_id = str(len(node_list))  # ID

        node_str = node_type + "\n" + node_text
        node_item = (node_id, node_str)  # tuple（node_id, node_string）  node_str是你想展示的内容
        node_list.append(node_item)

        # 如果有父节点，添加一条从父节点到当前节点的边
        if parent_id is not None:
            edge_list.append((parent_id, node_id))

        # 遍历子节点入队
        for child in node.children:
            queue.append((child, node_id))

    return node_list, edge_list


def draw_tree(graph_name: str, node_list: list[tuple], edge_list: list[tuple]):
    """ draw a tree
    """
    # 创建一个Digraph对象
    dot = Digraph(comment=graph_name)
    # 绘制结点
    for node in node_list:
        dot.node(name=node[0], label=node[1], shape="box")
    # 绘制连线
    for edge in edge_list:
        dot.edge(tail_name=edge[0], head_name=edge[1], lable="")
    return dot


print(cpp_code_snippet)
tree = cpp_parser.parse(bytes(cpp_code_snippet, "utf8"))
root_node = tree.root_node  # 注意，root_node 才是可遍历的树节点
node_list, edge_list = ast_bfs(root=root_node)  # BFS 遍历树
dot = draw_tree("test_graph", node_list, edge_list)
dot.view()

print('tree-sitter done.')

希望本篇文章对你有帮助，创作不易，感谢支持！