How to realize the visualization of syntax tree of simple four arithmetic operations in python

Simple four arithmetic operations syntax tree visualization

If you want to try it, you need to install this python visualization library.

Parsing AST

Python’s ast library has a parse method that can parse the incoming content into an AST. Then we export it using ast.dump and print it.
Note: indent This parameter is only available after Python 3.9. If the version is earlier, you can remove it, which will only affect the final output format.

Okay, it’s that simple. We've done that because this library is very powerful, but it's only used a little here. In fact, the basic structure can already be seen here, but my purpose is to generate a JSON representation of this tree. I want to use the above Python drawing library to draw it. The input it supports is JSON, and its format is:

{
	"name": "A",
	"children": [
		"name": "B",
		"children": []
	]
}

Rough traversal method

"""
Python's AST
利用 Python 的 ast 模块来解析表达式（简单的二元运算），
然后通过遍历语法树来生成 JSON 表示，再使用 PYthon 的库来
将其可视化。这个程序的目的是为了验证自己写的简易解析器是否正确。
"""

import ast
import json


# 操作类型和操作符映射的字典
OPERATORS = {
    ast.Add: "+",
    ast.Sub: "-",
    ast.Mult: "*",
    ast.Div: "/"
}


def generate(tree: ast.Module):
    """
    generate expression AST's representation of JSON
    """
    if not tree:
        raise Exception("Emtpy AST tree!")
    if tree.__class__ == ast.Module:
        print(json.dumps({
            "name": "Expr",
            "children": [DFS(tree.body[0].value)]  # type: ignore
        }, indent=4))


def DFS(node):
    """
    DFS AST
    """
    if not node:
        return {}

    if node.__class__ == ast.BinOp:
        return {
            "name": "BinOp",
            "children": [
                {
                    "name": "left",
                    "children": [
                        DFS(node.left)
                    ]
                },
                DFS(node.op),
                {
                    "name": "left",
                    "children": [
                        DFS(node.right)
                    ]
                }
            ]
        }

    if node.__class__ == ast.Constant:
        return {
            "name": "NUMBER",
            "children": [
                {
                    "name": str(node.value)  # Python 的绘图库，必须是字符串才能正常显示
                }
            ]
        }

    if node.__class__ in [ast.Add, ast.Sub, ast.Mult, ast.Div]:
        return {
            "name": "Op",
            "children": [
                {
                    "name": OPERATORS[node.__class__]
                }
            ]
        }

    # 这里我只处理 加减乘除和数字类型的运行
    raise Exception("There is not support extra type.")


if __name__ == "__main__":
    ast_tree = ast.parse("1+2+3+4+5")
    print(ast.dump(ast_tree, indent=4))
    generate(ast_tree)

Running result:
I will output two things here, one is the dump of the AST; the other is the JSON representation of the AST (the JSON representation of the logical structure, not the JSON representation of the object).

Rendering display

Copy the printed JSON string into the file and name it data.json. I think it's quite interesting to output directly to the console. I like to see the results directly.
Execute the following command: pytm-cli -d TB -i data.json -o demo.html
Open demo.html in the browser to see the effect .

Mainstream traversal methods

Although the above traversal method is easy to understand, it is difficult to expand. AST is usually traversed through Visitor pattern, and the ast library also provides several traversal methods.
Because we only need to traverse to generate JSON and do not need to modify the AST itself, we only look at the following two types. Obviously the first one cannot be used, and the reason has been marked in blue. It speaks for itself if you don't care about context, because generating JSON is actually all you need to care about. So, I choose ast.NodeVisitor below. Using it is also very simple. Just inherit this class and then write different processing logic for different nodes (this separates the logic of different nodes and reduces the coupling of the code).

Complete code

"""
Python's AST
利用 Python 的 ast 模块来解析表达式（简单的二元运算），
然后通过遍历语法树来生成 JSON 表示，再使用 PYthon 的库来
将其可视化。这个程序的目的是为了验证自己写的简易解析器是否正确。
"""

import ast
import json


# 操作类型和操作符映射的字典
OPERATORS = {
    ast.Add: "+",
    ast.Sub: "-",
    ast.Mult: "*",
    ast.Div: "/"
}

class JSONVisitor(ast.NodeVisitor):
    """
    JSON visitor: Traversal AST and generate JSON representation
    """

    def visit_Module(self, node):
        module = {
            "name": "Module",
            "children": []
        }
        for sub_node in node.body:
            module["children"].append(self.visit(sub_node))

        return module

    def visit_Expr(self, node):
        return {
            "name": "Expr",
            "children": [
                self.visit(node.value)
            ]
        }

    def visit_BinOp(self, node):
        return {
            "name": "BinOp",
            "children": [
                {
                    "name": "left",
                    "children": [
                        self.visit(node.left)
                    ]
                },
                self.visit(node.op),
                {
                    "name": "right",
                    "children": [
                        self.visit(node.right)
                    ]
                }
            ]
        }

    def visit_Constant(self, node):
        return {
            "name": "NUMBER",
            "children": [{
                "name": str(node.value)  # # Python 的绘图库，必须是字符串才能正常显示
            }]
        }

    def visit_Add(self, node):
        return self.__visit(node)

    def visit_Sub(self, node):
        return self.__visit(node)

    def visit_Mult(self, node):
        return self.__visit(node)

    def visit_Div(self, node):
        return self.__visit(node)

    def __visit(self, node):
        return {
            "name": "Op",
            "children": [{
                "name": OPERATORS[node.__class__]
            }]
        }


if __name__ == "__main__":
    ast_tree = ast.parse("1+2+3+4+5")
    visitor = JSONVisitor()
    json_str = visitor.visit(ast_tree)
    print(json.dumps(json_str, indent=4))

The previous rough version started directly from Expr. For this more elegant version, I will Module The node is also added.

The above is the detailed content of How to realize the visualization of syntax tree of simple four arithmetic operations in python. For more information, please follow other related articles on the PHP Chinese website!