AST way to evaluate an expression Step-by-step

ast_parse.py

import ast
from typing import Tuple, List
import black  # version: 22.10.0

BLACK_MODE = black.Mode(
    target_versions={black.TargetVersion.PY311}, line_length=120  # pyright:ignore
)
HORIZONTAL = "*"
SIDE = "*"
CORNER = "*"


def print_lines(lines: List[str]):
    max_length = 80
    for line in lines:
        if (length := len(line)) > max_length:
            max_length = length

    print(f"{CORNER}{HORIZONTAL * max_length}{CORNER}")
    for line in lines:
        print(f"{SIDE} {line}{' ' * (max_length - len(line) - 1)}{SIDE}")
    print(f"{CORNER}{HORIZONTAL * max_length}{CORNER}")


def invalid_syntax(expression: str = "", start: int = 0, end: int = 0):
    lines = [
        "OOPS! There is an error!",
        "Details:",
        expression,
        f"{' ' * start}^{'^' * (end - start)} SYNTAX ERROR",
        "Hint: Check if the input is correct.",
    ]
    print_lines(lines)


def invalid_symbol(expression: str = "", start: int = 0, end: int = 0):
    lines = [
        "OOPS! There is an error!",
        "Details:",
        expression,
        f"{' ' * start}{'^' * (end - start)} INVALID SYMBOL",
        "Hint: Check if the input is correct.",
    ]
    print_lines(lines)


def format_code(orig_code: str) -> str:
    """Formats code specified in orig_code and output the result."""
    code = orig_code
    try:
        code = black.format_file_contents(code, mode=BLACK_MODE, fast=True)
    except black.NothingChanged:  # pyright:ignore
        pass
    except black.InvalidInput:  # pyright:ignore
        invalid_syntax(orig_code)
        return "0"

    return code  # print result


def parse_expression(expression: str):
    """Parsess an expression and output the AST tree."""
    try:
        parsed_tree = ast.parse(expression, mode="eval")
    except SyntaxError:
        invalid_syntax(expression)
        exit(1)
    return parsed_tree


def get_binop(op):
    if isinstance(op, ast.Add):
        return "+"
    if isinstance(op, ast.Sub):
        return "-"
    if isinstance(op, ast.Mult):
        return "*"
    if isinstance(op, ast.Div):
        return "/"
    if isinstance(op, ast.FloorDiv):
        return "//"
    if isinstance(op, ast.Mod):
        return "%"
    if isinstance(op, ast.Pow):
        return "^"
    return


def get_unop(op):
    if isinstance(op, ast.UAdd):
        return "+"
    if isinstance(op, ast.USub):
        return "-"
    # NOTE: NOT and INVERT are **NOT** supported.
    return


def parse_part(
    orig_expression: str, value, ctx_orig: int = 0, no_print: bool = False
) -> Tuple[str, int]:
    ctx = ctx_orig
    expression = ""

    # Functions
    if isinstance(value, ast.Call):
        ctx += 1
        function_name = value.func.id  # pyright:ignore
        function_args = ", ".join(
            [parse_part(orig_expression, arg, ctx)[0] for arg in value.args]
        )  # pyright:ignore
        expression = f"{function_name}({function_args})"
        result_expression = (
            f"{orig_expression[:start]}{eval(expression)}{orig_expression[end:]}"
        )

    # -a...
    elif isinstance(value, ast.UnaryOp):
        ctx += 1
        op = get_unop(value.op)
        start = value.col_offset
        end = start + len(ast.unparse(value))
        if not op:
            invalid_symbol(orig_expression, start, end)

        operand = parse_part(orig_expression, value.operand, ctx)[0]
        expression = f"{op}{operand}"

        result_expression = (
            f"{orig_expression[:start]}{eval(expression)}{orig_expression[end:]}"
        )

    # a + b...
    elif isinstance(value, ast.BinOp):
        ctx += 1
        start = value.col_offset
        end = start + len(ast.unparse(value))
        left = parse_part(orig_expression, value.left, ctx)[0]
        right = parse_part(orig_expression, value.right, ctx)[0]

        op = get_binop(value.op)
        if not op:
            invalid_symbol(orig_expression, start, end)
            return
        expression = f"{left} {op} {right}"

        result_expression = (
            f"{orig_expression[:start]}{eval(expression)}{orig_expression[end:]}"
        )
        if ctx != ctx_orig:
            # Go deeper and check if we can still simplify the expression.
            print(f"= {result_expression}")
            if result_expression != orig_expression:
                parse_part(result_expression, parse_expression(result_expression), ctx)

    elif isinstance(value, ast.Constant):
        expression = str(value.value)

    elif isinstance(value, ast.Name):
        expression = value.id

    return expression, ctx


def evaluate(expression):
    expression = format_code(expression).strip()
    parsed_tree = parse_expression(expression)
    print(f"  {expression}")
    
    result = parse_part(expression, parsed_tree.body)


if __name__ == "__main__":
    # Replace with yours
    evaluate("1*2+3*1+3*2")

LICENSE

MIT License

Copyright (c) 2024 Andy Zhang

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Thanks!