simple evaluator with macros

main.py

from __future__ import annotations
import dataclasses
from typing import Any, TypeAlias

Expr: TypeAlias = Any


@dataclasses.dataclass
class Context:
    env: dict[str, Any] = dataclasses.field(default_factory=dict)

    def bind(self, vars: list[str], what: str) -> Context:
        return Context(self.env | {v: what for v in vars})


def evaluate(expr: Expr, ctx: Context) -> Any:
    return eval(translate_expr(expr, ctx))


def translate_expr(expr: Expr, ctx: Context) -> str:
    match expr:
        case list():
            return f"[{', '.join(translate_expr(x, ctx) for x in expr)}]"
        case int(x):
            return str(x)
        case str(s):
            return s
        case ("quote", value):
            return quote(value, ctx)
        case ("let-macro", ((name, *args), mbody), body):
            # here it becomes apparent that macros involve compile-time evaluation
            fun = evaluate(("lambda", args, mbody), ctx)
            _macros.add(fun)
            body_ctx = ctx.bind([name], fun)
            # no recursive macro expansion, yet
            return translate_expr(body, body_ctx)
        case ("+", *rands):
            return f"sum([{', '.join(translate_expr(r, ctx) for r in rands)}])"
        case ("lambda", params, body):
            body_ctx = ctx.bind(params, "VAR")
            return f"(lambda {', '.join(params)}: {translate_expr(body, body_ctx)})"
        case (rator, *rands) if is_macro(rator, ctx):
            macro = ctx.env[rator]
            expansion = macro(*rands)
            return translate_expr(expansion, ctx)
        case (rator, *rands):
            fun = translate_expr(rator, ctx)
            args = (translate_expr(a, ctx) for a in rands)
            return f"{fun}({', '.join(args)})"
        case _:
            raise TypeError(expr)


def quote(expr: Expr, ctx: Context) -> str:
    match expr:
        case list():
            return f"[{', '.join(quote(x, ctx) for x in expr)}]"
        case ("unquote", x):
            return translate_expr(x, ctx)
        case tuple():
            return f"({', '.join(quote(x, ctx) for x in expr)})"
        case _: return repr(expr)


def is_macro(expr: Expr, ctx: Context) -> bool:
    return isinstance(expr, str) and ctx.env.get(expr) in _macros


def gensym(name: str="") -> str:
    global _GENSYM_COUNTER
    _GENSYM_COUNTER += 1
    return f"___{name}{_GENSYM_COUNTER}"
_GENSYM_COUNTER = 0

_macros = set()
default_ctx = Context()


# an example macro, written in Python
def let_macro(defs: Expr, body: Expr) -> Expr:
    vars = [d[0] for d in defs]
    vals = [d[1] for d in defs]
    return (("lambda", vars, body), *vals)


_macros.add(let_macro)
default_ctx.env["let"] = let_macro


# an example macro, written in Python
def begin_macro(*args: Expr) -> Expr:
    vars = [gensym() for _ in args]
    return (("lambda", vars, vars[-1]), *args)


_macros.add(begin_macro)
default_ctx.env["begin"] = begin_macro


def play(expr: Expr):
    print(expr)
    print("   ", translate_expr(expr, default_ctx))
    print("   ", evaluate(expr, default_ctx))


play(("quote", (("lambda", ["x"], "x"), 42)))
play((("lambda", ["x"], "x"), 42))
play(("let", (("x", 123), ("y", 42)), ("+", "x", "y")))
play(("begin", 1, 2, 3))
play(("begin", ("begin", 1, 2, 3), ("begin", 4, 5, 6)))
play(("let-macro", (("inc", "n"), ("quote", ("+", 1, ("unquote", "n")))), ("inc", ("inc", 40))))