如何检索python代码的控制流图？

import ast
import graphviz

def build_control_flow_graph(code):
    # 解析代码为抽象语法树
    tree = ast.parse(code)

    # 构建控制流图的数据结构
    graph = graphviz.Digraph()

    # 遍历抽象语法树
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            # 添加函数节点
            graph.node(node.name, shape='box')
            # 添加函数内部的控制流关系
            for child_node in ast.iter_child_nodes(node):
                if isinstance(child_node, ast.If):
                    # 添加条件语句节点
                    graph.node(str(child_node.lineno), shape='diamond')
                    # 添加条件语句与函数之间的控制流关系
                    graph.edge(node.name, str(child_node.lineno))
                    # 添加条件语句内部的控制流关系
                    for if_child_node in ast.iter_child_nodes(child_node):
                        if isinstance(if_child_node, ast.Return):
                            # 添加返回语句节点
                            graph.node(str(if_child_node.lineno), shape='box')
                            # 添加返回语句与条件语句之间的控制流关系
                            graph.edge(str(child_node.lineno), str(if_child_node.lineno))
                elif isinstance(child_node, ast.Return):
                    # 添加返回语句节点
                    graph.node(str(child_node.lineno), shape='box')
                    # 添加返回语句与函数之间的控制流关系
                    graph.edge(node.name, str(child_node.lineno))

    return graph

# 示例代码
code = '''
def foo(x):
    if x > 0:
        return x
    else:
        return -x
'''

# 构建控制流图
graph = build_control_flow_graph(code)

# 可视化控制流图
graph.render('control_flow_graph', format='png')