四則演算 (1桁の数値) の数式からASTを生成

graphviz.py

""" Graphviz で使えるテキストを吐き出すやつ
"""

from node import Node, NodeKind


def gen(node: Node, node_id=""):
    if node.kind == NodeKind.NUM:
        # 数値は値を表示
        print(f'\t{node_id} [label="{node.value}"];')
        return

    gen(node.left, node_id + "l")
    gen(node.right, node_id + "r")

    label = ""
    if node.kind == NodeKind.ADD:
        label = "+"
    elif node.kind == NodeKind.SUB:
        label = "-"
    elif node.kind == NodeKind.MUL:
        label = "*"
    elif node.kind == NodeKind.DIV:
        label = "/"

    print(f'\t{node_id} [label="{label}"];')
    print(f'\t{node_id} -> {node_id + "l"};')
    print(f'\t{node_id} -> {node_id + "r"};')


def generate_dot(tree: Node, name: str = "sample"):
    """ Graphviz でそのまま出力できるテキスト"""
    print("digraph sample {")
    # print("\tnode [ shape = circle];")
    gen(tree, 'root')
    print("}")

node.py

from __future__ import annotations
from enum import Enum, auto
from typing import Optional


class NodeKind(Enum):
    """ 節点の種類 """

    ADD = auto()  # 加算 Add
    SUB = auto()  # 減算 Subtract
    MUL = auto()  # 乗算 Multiply
    DIV = auto()  # 除算 Division
    NUM = auto()


class Node:
    """ 節点を表す """

    def __init__(
        self,
        kind: NodeKind,
        left: Optional[Node],
        right: Optional[Node],
        value: int = 0,
    ) -> None:
        self.kind: NodeKind = kind
        # 数値の場合、値が入る
        self.value: int = value
        self.left: Optional[Node] = left
        self.right: Optional[Node] = right

pparser_ast.py

"""
中置記法の数式のASTを作成する

graphviz のテキストも出力できる

expr   -> expr + term
        | expr - term
        | term
term   -> term * factor
        | term / factor
        | factor
factor -> digit
        | ( expr )

意味規則を追加する (抽象構文木を生成する)
    expr の属性n に節点を付与 (実際には return ?)
    解析結果を節点の子として、設定するため、一時変数に保持する
    (- と / は省略)

expr   -> expr  { e = expr() }
          +     { op = '+' }
          term  { t = term()
                  expr.n = new Node(op, left=e, right=t) }
        | term  { expr.n = term() }
term   -> term  { e = term() }
          *     { op = '*' }
          term  { t = term()
                  term.n = new Node(op, left=e, right=t) }
        | factor { term.n = factor() }
factor -> digit { factor.n = new Node(t=NUM, l=null, r=null) }
        | (  expr { factor.n = expr() } )


の左再帰を取り除く
    左側の子を作っていくイメージ？
        E や T には、左の子となる部分木を渡す
            E, T の処理ではそれを使って、部分木を作成し、再帰する
        節点を作るための子になるため
            左の子にするのは、左結合だからかな？
    空のときは、そのまま返す

expr   -> term  { left = term() }
          E     { expr.n = E(left) }
E      -> +     { op = '+' }
          E     { right = term()
                  E.n = E(new Node(op, left, right)) }
        | e     { E.n = left }
term   -> factor { left = factor() }
          T      { term.n = T(left) }
T      -> *      { op = '*' }
          factor { right = factor() }
          T      { T.n = T(new Node(op, left, right)) }
        | e     { T.n = left }
factor -> digit { factor.n = new Node(t=NUM, l=null, r=null) }
        | (  expr { factor.n = expr() } )


"""
import string
from typing import Optional

from graphviz import generate_dot
from node import Node, NodeKind


def is_digit(s) -> bool:
    return s in string.digits


class Parser:
    def __init__(self, input_str):
        self.input_str = input_str
        self.pos = 0

    @property
    def lookahead(self) -> Optional[str]:
        if self.pos >= len(self.input_str):
            return None
        return self.input_str[self.pos]

    def next(self) -> Optional[str]:
        self.pos += 1
        return self.lookahead

    def match(self, s) -> None:
        if self.lookahead and self.lookahead == s:
            self.next()
        else:
            raise Exception("syntax error")

    def parse(self) -> Node:
        return self.expr()

    def expr(self) -> Node:
        left = self.term()
        return self.E(left)

    def E(self, left=None) -> Node:
        if self.lookahead and self.lookahead == "+":
            self.match("+")
            right = self.term()
            # 節点を作り、左の部分木としてつなげてもらう
            n = Node(NodeKind.ADD, left, right)
            return self.E(n)
        elif self.lookahead and self.lookahead == "-":
            self.match("-")
            right = self.term()
            n = Node(NodeKind.SUB, left, right)
            return self.E(n)
        else:
            # 空列
            # そのまま返す
            return left

    def term(self) -> Node:
        left = self.factor()
        return self.T(left)

    def T(self, left) -> Node:
        if self.lookahead and self.lookahead == "*":
            self.match("*")
            right = self.factor()
            n = Node(NodeKind.MUL, left, right)
            return self.T(n)
        elif self.lookahead and self.lookahead == "/":
            self.match("/")
            right = self.factor()
            n = Node(NodeKind.DIV, left, right)
            return self.T(n)
        else:
            # 空列
            # そのまま返す
            return left

    def factor(self) -> Node:
        if self.lookahead and is_digit(self.lookahead):
            n = self.lookahead
            self.match(self.lookahead)
            return Node(NodeKind.NUM, None, None, int(n))
        elif self.lookahead and self.lookahead == "(":
            self.match("(")
            ret = self.expr()
            self.match(")")
            return ret
        else:
            raise Exception("syntax error")


def main():
    parser = Parser("1+2+3*3+1")
    tree = parser.parse()

    generate_dot(tree)


if __name__ == "__main__":
    main()

これ、正しいかも...？