dc_parser.py

from dataclasses import dataclass
from enum import Enum
import operator
from typing import Optional


class TokenType(Enum):
    LPAREN = 1
    RPAREN = 2
    PLUS = 3
    MINUS = 4
    MULTIPLY = 5
    DIVIDE = 6
    POWER = 7
    INT = 8
    FLOAT = 9


class NodeType(Enum):
    ADD = operator.add
    SUB = operator.sub
    MUL = operator.mul
    DIV = operator.truediv
    POW = operator.pow
    NEG = operator.neg
    INT = int
    FLOAT = float


ops = {
    "ADD": "+",
    "SUB": "-",
    "MUL": "*",
    "DIV": "/",
    "POW": "^",
    "NEG": "-",
}


class Token:
    def __init__(self, type, value=None):
        self.type = type
        self.value = value


@dataclass(init=False)
class Lexer:
    text: str
    pos: int
    cur_char: Optional[str]

    def __init__(self, text: str, pos: int = 0):
        self.text = text
        self.pos = pos
        self.cur_char = text[pos] if pos < len(text) else None

    tokens = {
        "(": TokenType.LPAREN,
        ")": TokenType.RPAREN,
        "+": TokenType.PLUS,
        "-": TokenType.MINUS,
        "*": TokenType.MULTIPLY,
        "/": TokenType.DIVIDE,
        "^": TokenType.POWER,
    }

    def step(self):
        self.pos += 1
        if self.pos > len(self.text) - 1:
            self.cur_char = None
        else:
            self.cur_char = self.text[self.pos]
        return self.cur_char

    def get_next_token(self):
        while self.cur_char is not None:
            if self.cur_char.isspace():
                self.step()
                continue

            if self.cur_char in self.tokens:
                token_type = self.tokens[self.cur_char]
                self.step()
                return Token(token_type, self.text[self.pos - 1])

            if self.cur_char.isdigit() or self.cur_char == ".":
                num_str = ""
                is_float = False

                while self.cur_char is not None and (
                    self.cur_char.isdigit() or self.cur_char == "."
                ):
                    if self.cur_char == ".":
                        is_float = True
                    num_str += self.cur_char
                    self.step()

                try:
                    if not is_float:
                        return Token(TokenType.INT, int(num_str))
                    else:
                        return Token(TokenType.FLOAT, float(num_str))
                except ValueError:
                    raise Exception(
                        f"Invalid number: {num_str}, at position {self.pos + 1}"
                    )

            raise Exception(
                f"Invalid character: {self.cur_char}, at position {self.pos + 1}"
            )

        return None

    def __iter__(self):
        return self

    def __next__(self):
        tok = self.get_next_token()
        if tok is None:
            raise StopIteration
        return tok


@dataclass(slots=True)
class Branch:
    nodetype: NodeType
    left: "Branch | Leaf"
    right: "Branch | Leaf"

    def __repr__(self):
        return f"{self.nodetype.name}({self.left}, {self.right})"

    def __str__(self):
        return f"({self.left}{ops[self.nodetype.name]}{self.right})"

    def __call__(self):
        return self.nodetype.value(self.left(), self.right())


@dataclass(slots=True)
class Leaf:
    nodetype: NodeType
    value: int | float | None = None

    def __repr__(self):
        return f"{self.nodetype.type.name}({self.value})"

    def __str__(self):
        return f"{'-' if self.nodetype == NodeType.NEG else ''}{self.value}"

    def __call__(self):
        return self.nodetype.value(
            self.value
            if self.nodetype in (NodeType.INT, NodeType.FLOAT)
            else self.value()
        )


class Parser:
    def __init__(self, lexer):
        self.lexer = lexer
        self._iter = iter(self.lexer)
        self.current_token = self.step()

    def step(self):
        try:
            self.current_token = next(self.lexer)
            return self.current_token
        except StopIteration:
            self.current_token = None
            return None

    def consume(self, token_type):
        if self.current_token is None:
            raise SyntaxError(
                f"Expected token {token_type.name}, but reached end of input at position {self.lexer.pos + 1}\n"
                f"{self.lexer.text}\n{' ' * (self.lexer.pos-1)}^"
            )
        if self.current_token.type == token_type:
            self.step()
        else:
            raise SyntaxError(
                f"Expected token {token_type.name}, got {self.current_token.type.name} at position {self.lexer.pos + 1}\n"
                f"{self.lexer.text}\n{' ' * (self.lexer.pos-1)}^"
            )

    def parse(self):
        return self.expr()

    def expr(self):
        node = self.term()
        token = self.current_token
        if token is not None:
            if token.type == TokenType.PLUS:
                self.step()
                node = Branch(NodeType.ADD, node, self.expr())
            elif token.type == TokenType.MINUS:
                self.step()
                node = Branch(NodeType.SUB, node, self.expr())
        return node

    def term(self):
        node = self.factor()
        token = self.current_token
        if token is not None:
            if token.type == TokenType.MULTIPLY:
                self.step()
                node = Branch(NodeType.MUL, node, self.term())
            elif token.type == TokenType.DIVIDE:
                self.step()
                node = Branch(NodeType.DIV, node, self.term())
        return node

    def factor(self):
        node = self.atom()
        token = self.current_token
        if token is not None and token.type == TokenType.POWER:
            self.step()
            node = Branch(NodeType.POW, node, self.factor())
        return node

    def atom(self):
        if self.current_token is None:
            raise SyntaxError(
                f"Unexpected end of input at position {self.lexer.pos + 1}\n"
                f"{self.lexer.text}\n{' ' * (self.lexer.pos-1)}^"
            )

        match self.current_token.type:
            case TokenType.PLUS:
                self.step()
                return self.factor()
            case TokenType.MINUS:
                self.step()
                return Leaf(NodeType.NEG, self.factor())
            case TokenType.INT:
                value = self.current_token.value
                self.step()
                return Leaf(NodeType.INT, value)
            case TokenType.FLOAT:
                value = self.current_token.value
                self.step()
                return Leaf(NodeType.FLOAT, value)
            case TokenType.LPAREN:
                self.consume(TokenType.LPAREN)
                node = self.expr()
                self.consume(TokenType.RPAREN)
                return node
            case _:
                raise SyntaxError(
                    f"Unexpected token {self.current_token.type.name} at position {self.lexer.pos + 1}\n"
                    f"{self.lexer.text}\n{' ' * (self.lexer.pos-1)}^"
                )


if __name__ == "__main__":
    while True:
        input_string = input("Enter an expression: ")
        lexer = Lexer(input_string)
        parser = Parser(lexer)
        equation = parser.parse()
        answer = equation()
        print(f"Result: {equation} = {answer}")