python macros

letfy.py

from ast import *


class LetVisitor(NodeTransformer):
    def visit_Call(self, node):

        self.generic_visit(node)

        if (
            type(node.func) is Call
            and type(node.func.func) is Name
            and node.func.func.id == "let"
        ):
            node.args = [Constant(unparse(node.args[0]))]

        return node

    def visit_FunctionDef(self, node):
        "Removes the decorater to avoid infinite loop"
        found = None
        node.decorator_list = list(
            filter(lambda d: d.id != "letfy", node.decorator_list)
        )
        self.generic_visit(node)
        return node


def letfy(f):
    import inspect, types

    source = inspect.getsource(f)
    print("original source\n")
    print(source, "\n\n")
    old_code_obj = f.__code__
    old_ast = parse(source)
    new_ast = fix_missing_locations(LetVisitor().visit(old_ast))
    new_code_obj = compile(new_ast, old_code_obj.co_filename, "exec")
    print("new source\n")
    print(unparse(new_ast), "\n\n")
    new_f = types.FunctionType(new_code_obj.co_consts[0], f.__globals__)
    return new_f


def let(**kwargs):
    def _body(expr):
        keys = ",".join(kwargs.keys())
        kvals = ",".join(f"{k}={v}" for k, v in kwargs.items())
        return eval(f"(lambda {keys}: {expr})({kvals})")

    return _body


@letfy
def foo():
    return let(a=1, b=2)(a + b)


print(foo())