Lisp but none of the fun parts

lisp.py

from dataclasses import dataclass
from typing import Literal

type LeftParen = Literal["("]
type RightParen = Literal[")"]
type Number = int
type Token = LeftParen | RightParen | NumberLiteral | Identifier | Def | Eof | If | StringLiteral | LogicLiteral


@dataclass
class NumberLiteral:
    number: int


@dataclass
class StringLiteral:
    value: str


@dataclass
class Identifier:
    name: str


@dataclass
class LogicLiteral:
    value: bool


@dataclass
class Def:
    pass


@dataclass
class If:
    pass


@dataclass
class Eof:
    pass


def isident(c: str) -> bool:
    if c.isspace():
        return False
    for c in c:
        if c in "()[]{}\",'`;#|\\":
            return False
    return True


def tokenize(src: str) -> list[Token]:
    idx = 0
    tokens: list[Token] = []
    while idx < len(src):
        c = src[idx]
        if c == "(" or c == ")":
            tokens.append(c)
            idx += 1
        elif c == '"':
            idx += 1
            string = ""
            escaping = False
            while idx < len(src):
                if escaping and src[idx] == "n":
                    string += "\n"
                    escaping = False
                    idx += 1
                elif escaping and src[idx] == "\\":
                    string += "\\"
                    escaping = False
                    idx += 1
                elif escaping and src[idx] == '"':
                    string += '"'
                    escaping = False
                    idx += 1
                elif escaping:
                    assert False, "Unknown escape sequence"
                elif src[idx] == '"':
                    break
                else:
                    string += src[idx]
                    idx += 1
            assert idx < len(src) and src[idx] == '"', "Unclosed string"
            idx += 1
            tokens.append(StringLiteral(string))
        elif isident(c):
            ident = ""
            while idx < len(src) and isident(src[idx]):
                ident += src[idx]
                idx += 1
            if ident.isdigit():
                tokens.append(NumberLiteral(int(ident)))
            elif ident == "def":
                tokens.append(Def())
            elif ident == "if":
                tokens.append(If())
            elif ident == "true":
                tokens.append(LogicLiteral(True))
            elif ident == "false":
                tokens.append(LogicLiteral(False))
            else:
                tokens.append(Identifier(ident))
        else:
            idx += 1

    tokens.append(Eof())
    return tokens


"""
    program ::= expr* EOF

    expr ::= def | call | if | IDENTIFIER | NUMBER | STRING | BOOLEAN

    call ::= "(" IDENTIFIER expr* ")"

    if ::= "(" "if" expr expr expr ")"

    def ::= "(" "def" "(" IDENTIFIER+ ")" expr+ ")"
"""

type Program = list[Expr]
type Expr = NumberExpr | StringExpr | LogicExpr | VarExpr | CallExpr | DefExpr | IfExpr


@dataclass
class NumberExpr:
    number: int


@dataclass
class LogicExpr:
    value: bool


@dataclass
class StringExpr:
    value: str


@dataclass
class VarExpr:
    var: str


@dataclass
class CallExpr:
    fname: str
    args: list[Expr]


@dataclass
class DefExpr:
    iname: str
    params: list[Identifier]
    exprs: list[Expr]


@dataclass
class IfExpr:
    cond: Expr
    then_b: Expr
    else_b: Expr


def expr(tokens: list[Token]) -> tuple[Expr, list[Token]]:
    """returns an expr and how many tokens were consumed"""
    first, *rest = tokens

    if isinstance(first, NumberLiteral):
        return NumberExpr(first.number), rest
    elif isinstance(first, StringLiteral):
        return StringExpr(first.value), rest
    elif isinstance(first, Identifier):
        return VarExpr(first.name), rest

    if first == "(":
        first, *rest = rest
        match first:
            case If():
                cond, rest = expr(rest)
                thenb, rest = expr(rest)
                elseb, rest = expr(rest)

                # consume right bracket
                assert rest[0] == ")", "Unclosed def"
                _, *rest = rest

                return IfExpr(cond, thenb, elseb), rest
            case Def():
                # consume def, get other
                first, *rest = rest
                assert first == "(", "Token after 'def' must be left parens"

                iname, *rest = rest
                assert isinstance(
                    iname, Identifier
                ), "Definitions must have at least a name"

                params: list[Identifier] = []
                while rest[0] != ")":
                    arg, *rest = rest
                    assert isinstance(arg, Identifier), "Parameters must be identifiers"
                    params.append(arg)

                first, *rest = rest
                assert first == ")", "Parameter list must be closed"

                exprs: list[Expr] = []
                while rest[0] != ")":
                    e, rest = expr(rest)
                    exprs.append(e)

                # consume right bracket
                assert rest[0] == ")", "Unclosed def"
                _, *rest = rest

                return DefExpr(iname.name, params, exprs), rest
            case Identifier(fname):
                args: list[Expr] = []
                while rest[0] != ")":
                    e, rest = expr(rest)
                    args.append(e)

                assert rest[0] == ")", "Unclosed call"
                _, *rest = rest

                return CallExpr(fname, args), rest
            case _:
                pass

    assert False, "Expected expression"


def reader(src: str) -> Program:
    tokens = tokenize(src)
    program: Program = []
    i = 0
    while i < len(tokens):
        token = tokens[i]
        if isinstance(token, Eof):
            i += 1
            break

        exp, tokens = expr(tokens)
        program.append(exp)

    return program


type Value = IntValue | StringValue | FunctionValue | LogicValue | BuiltinFunctionValue


@dataclass
class IntValue:
    value: int


@dataclass
class StringValue:
    value: str


@dataclass
class LogicValue:
    value: bool


@dataclass
class FunctionValue:
    name: str
    params: list[str]
    body: list[Expr]


@dataclass
class BuiltinFunctionValue:
    name: str
    params: list[str]


def value_tostring(value: Value) -> str:
    match value:
        case IntValue(v):
            return str(v)
        case StringValue(v):
            return v
        case LogicValue(v):
            return "true" if v else "false"
        case FunctionValue(name, _, _):
            return f"<function '{name}'@{hex(id(value))}>"
        case BuiltinFunctionValue(name, _):
            return f"<built-in function '{name}'>"


builtins = {
    "print": BuiltinFunctionValue(name="print", params=["arg"]),
    "read": BuiltinFunctionValue(name="read", params=[]),
    "exit": BuiltinFunctionValue(name="exit", params=[]),
    "<=": BuiltinFunctionValue(name="<=", params=["lhs", "rhs"]),
    ">=": BuiltinFunctionValue(name=">=", params=["lhs", "rhs"]),
    "<": BuiltinFunctionValue(name="<", params=["lhs", "rhs"]),
    ">": BuiltinFunctionValue(name=">", params=["lhs", "rhs"]),
    "=": BuiltinFunctionValue(name="=", params=["lhs", "rhs"]),
    "not": BuiltinFunctionValue(name="not", params=["arg"]),
    "+": BuiltinFunctionValue(name="+", params=["lhs", "rhs"]),
    "-": BuiltinFunctionValue(name="-", params=["lhs", "rhs"]),
    "*": BuiltinFunctionValue(name="*", params=["lhs", "rhs"]),
    "/": BuiltinFunctionValue(name="/", params=["lhs", "rhs"]),
}


def builtin_eval(
    f: BuiltinFunctionValue, values: list[Value], environ: dict[str, Value]
) -> Value:
    match f.name:
        case "print":
            print(value_tostring(values[0]))
            return LogicValue(True)
        case "read":
            return IntValue(int(input()))
        case "exit":
            exit(0)
        case "<=":
            assert isinstance(values[0], IntValue), "Arguments to <= must be integers"
            assert isinstance(values[1], IntValue), "Arguments to <= must be integers"
            return LogicValue(values[0].value <= values[1].value)
        case ">=":
            assert isinstance(values[0], IntValue), "Arguments to >= must be integers"
            assert isinstance(values[1], IntValue), "Arguments to >= must be integers"
            return LogicValue(values[0].value >= values[1].value)
        case "<":
            assert isinstance(values[0], IntValue), "Arguments to < must be integers"
            assert isinstance(values[1], IntValue), "Arguments to < must be integers"
            return LogicValue(values[0].value < values[1].value)
        case ">":
            assert isinstance(values[0], IntValue), "Arguments to > must be integers"
            assert isinstance(values[1], IntValue), "Arguments to > must be integers"
            return LogicValue(values[0].value > values[1].value)
        case "=":
            match values[0], values[1]:
                case IntValue(l), IntValue(r):
                    return LogicValue(l == r)
                case StringValue(l), StringValue(r):
                    return LogicValue(l == r)
                case _:
                    assert False, "Arguments to = must be of the same type"
        case "not":
            assert isinstance(
                values[0], LogicValue
            ), "Arguments to not must be booleans"
            return LogicValue(not values[0].value)
        case "+":
            assert isinstance(values[0], IntValue), "Arguments to + must be integers"
            assert isinstance(values[1], IntValue), "Arguments to + must be integers"
            return IntValue(values[0].value + values[1].value)
        case "-":
            assert isinstance(values[0], IntValue), "Arguments to - must be integers"
            assert isinstance(values[1], IntValue), "Arguments to - must be integers"
            return IntValue(values[0].value - values[1].value)
        case "*":
            assert isinstance(values[0], IntValue), "Arguments to * must be integers"
            assert isinstance(values[1], IntValue), "Arguments to * must be integers"
            return IntValue(values[0].value * values[1].value)
        case "/":
            assert isinstance(values[0], IntValue), "Arguments to / must be integers"
            assert isinstance(values[1], IntValue), "Arguments to / must be integers"
            return IntValue(values[0].value // values[1].value)
        case _:
            assert False, "Unknown built-in function"


def evaluate(expr: Expr, environ: dict[str, Value]) -> Value:
    match expr:
        case NumberExpr(x):
            return IntValue(x)
        case StringExpr(x):
            return StringValue(x)
        case LogicExpr(x):
            return LogicValue(x)
        case VarExpr(name):
            assert name in environ or name in builtins, f"'{name}' is not declared"
            return environ[name] if name in environ else builtins[name]
        case IfExpr(cond, thenb, elseb):
            v = evaluate(cond, environ)
            if isinstance(v, LogicValue) and v.value:
                return evaluate(thenb, environ)
            else:
                return evaluate(elseb, environ)
        case CallExpr(name, args):
            # get function
            f = (
                environ[name]
                if name in environ
                else builtins[name] if name in builtins else None
            )
            assert f, f"'{name}' is not declared"
            assert isinstance(f, FunctionValue) or isinstance(
                f, BuiltinFunctionValue
            ), f"'{name}' is not a function"

            # get parameter values (eager evaluation)
            values = [evaluate(arg, environ) for arg in args]
            assert len(values) == len(
                f.params
            ), f"Wrong arity for '{name}': expected {len(f.params)}, got {len(values)}"

            if isinstance(f, BuiltinFunctionValue):
                return builtin_eval(f, values, environ)

            # construct environment
            bindings = dict(zip(f.params, values))
            new_env = {**environ, **bindings}

            # evaluate
            assert f.body, "Function body is empty"
            for exp in f.body:
                v = evaluate(exp, new_env)

            return v  # type: ignore
        case DefExpr(name, params, exprs):
            # construct FunctionValue
            fv = FunctionValue(name, [p.name for p in params], exprs)
            environ[name] = fv

            return fv


def run_file(path: str):
    with open(path, mode="r", encoding="utf-8") as f:
        src = f.read()
    prog = reader(src)
    environ = {}
    for ex in prog:
        evaluate(ex, environ)


def repl():
    environ = {}
    while True:
        try:
            src = input("lisp> ")
        except KeyboardInterrupt:
            break

        if src == "exit":
            break

        try:
            prog = reader(src)
            for ex in prog:
                v = evaluate(ex, environ)
                print(value_tostring(v))
        except Exception as e:
            print(f"! Error: {e}")


if __name__ == "__main__":
    import sys

    match sys.argv[1:]:
        case []:
            repl()
        case [file]:
            run_file(file)
        case _:
            print("Usage: python3 lisp.py [file]")
            sys.exit(1)