prefix expression parser

__main__.py

#!/usr/bin/env python3

import sys
import re
from enum import auto, IntEnum
from typing import Any


re_number = re.compile(r'[0-9]+(.[0-9]+)?')


class Token:
    class Kind(IntEnum):
        INVALID = auto()
        NEW_LINE = auto()
        PAREN_OPEN = auto()
        PAREN_CLOSE = auto()
        SYMBOL = auto()
        OPERATOR_ADD = auto()
        OPERATOR_SUB = auto()
        OPERATOR_MUL = auto()
        OPERATOR_DIV = auto()

    __slots__ = ('kind', 'text')

    def __init__(self, kind: Kind = Kind.INVALID, text: str = '') -> None:
        self.kind = kind
        self.text = text

    def __repr__(self) -> str:
        return f'(kind: {self.kind.name}, "{self.text}")'


class TokenStream:
    def __init__(self, source_text: str) -> None:
        self.tokens: list[Token] = []

        i = 0

        while True:
            token = Token()

            if source_text[i].isspace():
                if source_text[i] == '\n':
                    token.kind = Token.Kind.NEW_LINE
                    self.tokens.append(token)
                    i += 1

                while source_text[i].isspace():
                    i += 1

                    if i >= len(source_text):
                        break
            else:
                tbl = {
                    '(': Token.Kind.PAREN_OPEN,
                    ')': Token.Kind.PAREN_CLOSE,
                    '+': Token.Kind.OPERATOR_ADD,
                    '-': Token.Kind.OPERATOR_SUB,
                    '*': Token.Kind.OPERATOR_MUL,
                    '/': Token.Kind.OPERATOR_DIV,
                }

                token.kind = Token.Kind.SYMBOL

                if source_text[i] in tbl:
                    token.kind = tbl[source_text[i]]
                    token.text = source_text[i]
                    i += 1
                else:
                    while not source_text[i].isspace() and source_text[i] not in tbl:
                        token.text += source_text[i]
                        i += 1

                        if i >= len(source_text):
                            break

                self.tokens.append(token)

            if i >= len(source_text):
                break


class Parser:
    def __init__(self, arg) -> None:
        self.i = 0

        if arg is TokenStream:
            self.token_stream = arg
        else:
            self.token_stream = TokenStream(arg)


    def parse_add(self):
        node: dict[str, Any] = {
            'kind': 'expr',
            'subkind': 'add',
        }

        self.i += 1 # skip +

        if self.token_stream.tokens[self.i].kind != Token.Kind.PAREN_OPEN:
            raise Exception(f"Expected `(', got `{self.token_stream.tokens[self.i].text}'.")

        self.i += 1 # skip (
        node['lhs'] = self.parse_expr()
        node['rhs'] = self.parse_expr()
        self.i += 1 # skip )

        return node


    def parse_mul(self):
        node: dict[str, Any] = {
            'kind': 'expr',
            'subkind': 'mul',
        }

        self.i += 1 # skip +

        if self.token_stream.tokens[self.i].kind != Token.Kind.PAREN_OPEN:
            raise Exception(f"Expected `(', got `{self.token_stream.tokens[self.i].text}'.")

        self.i += 1 # skip (
        node['lhs'] = self.parse_expr()
        node['rhs'] = self.parse_expr()
        self.i += 1 # skip )

        return node


    def parse_div(self):
        node: dict[str, Any] = {
            'kind': 'expr',
            'subkind': 'div',
        }

        self.i += 1 # skip +

        if self.token_stream.tokens[self.i].kind != Token.Kind.PAREN_OPEN:
            raise Exception(f"Expected `(', got `{self.token_stream.tokens[self.i].text}'.")

        self.i += 1 # skip (
        node['lhs'] = self.parse_expr()
        node['rhs'] = self.parse_expr()
        self.i += 1 # skip )

        return node


    def parse_sub(self):
        node: dict[str, Any] = {
            'kind': 'expr',
            'subkind': 'sub',
        }

        self.i += 1 # skip +

        if self.token_stream.tokens[self.i].kind != Token.Kind.PAREN_OPEN:
            raise Exception(f"Expected `(', got `{self.token_stream.tokens[self.i].text}'.")

        self.i += 1 # skip (
        node['lhs'] = self.parse_expr()
        node['rhs'] = self.parse_expr()
        self.i += 1 # skip )

        return node


    def parse_expr(self) -> dict:
        node = {}

        if self.token_stream.tokens[self.i].kind == Token.Kind.OPERATOR_ADD:
            node = self.parse_add()
        elif self.token_stream.tokens[self.i].kind == Token.Kind.OPERATOR_SUB:
            node = self.parse_sub()
        elif self.token_stream.tokens[self.i].kind == Token.Kind.OPERATOR_MUL:
            node = self.parse_mul()
        elif self.token_stream.tokens[self.i].kind == Token.Kind.OPERATOR_DIV:
            node = self.parse_div()
        elif re_number.match(self.token_stream.tokens[self.i].text):
            node['kind'] = 'numeric literal'
            node['value'] = float(self.token_stream.tokens[self.i].text)
            self.i += 1

        return node


def evaluate(ast):
    if ast['kind'] == 'expr':
        lhs = evaluate(ast['lhs'])
        rhs = evaluate(ast['rhs'])

        if ast['subkind'] == 'add':
            return lhs + rhs
        elif ast['subkind'] == 'sub':
            return lhs - rhs
        elif ast['subkind'] == 'mul':
            return lhs * rhs
        elif ast['subkind'] == 'div':
            return lhs / rhs
    elif ast['kind'] == 'numeric literal':
        return ast['value']


if __name__ == '__main__':
    print(evaluate(Parser(sys.argv[1]).parse_expr()))