Magic Python, for fun and 0 profit

mpy.py

import ast

_FILENAME = "<undefined>"


def magic(source, filename=_FILENAME):
    a = ast.parse(source)
    return Magic(filename).visit(a)


def macro(f):
    f._macro = True
    return f


def _is_macro(fun):
    return getattr(fun, "_macro", False)


def _apply_macro(fun, node, *args):
    return ast.fix_missing_locations(ast.copy_location(fun(*args), node))


class Magic(ast.NodeTransformer):
    def __init__(self, filename, ns=None):
        self.ns = ns or {}
        self.filename = filename

    def visit_Call(self, node):
        node = self.generic_visit(node)
        fun = self.ns.get(getattr(node.func, "id", None))
        if _is_macro(fun):
            node = _apply_macro(fun, node, *node.args)

        return node

    def _visit_def(self, node):
        if node.decorator_list:
            fun = self.ns.get(getattr(node.decorator_list[0], "id", None))
            if _is_macro(fun):
                node = _apply_macro(fun, node, node)

        return node

    def visit_ClassDef(self, node):
        node = self.generic_visit(node)
        node = self._visit_def(node)
        return node

    def visit_FunctionDef(self, node):
        node = self.generic_visit(node)
        node = self._visit_def(node)
        self.exec(node)
        return node

    def visit_Import(self, node):
        node = self.generic_visit(node)
        self.exec(node)
        return node

    def visit_ImportFrom(self, node):
        node = self.generic_visit(node)
        self.exec(node)
        return node

    def exec(self, node):
        exec(compile(ast.Module(body=[node]), self.filename, "exec"), self.ns)

    #def visit_Module(self, node):
    #    return ast.Module(body=[
    #        *node.body,
    #        ast.Assign(
    #            targets=[ast.Name(id="__macros__", ctx=ast.Store())],
    #            value=ast.List(elts=[ast.Constant(value=n) for n in self.ns],
    #                           ctx=ast.Load()))])

Example

In [91]: a = mpy.magic("""
    ...: import ast
    ...: 
    ...: from mpy import macro
    ...: 
    ...: @macro
    ...: def foo(a):
    ...:     print("got", ast.dump(a))
    ...:     return ast.Starred(value=ast.List(
    ...:         elts=[ast.Constant(value="HELLO WORLD"),
    ...:               ast.BinOp(left=a.left, op=ast.Sub(), right=a.right)],
    ...:         ctx=ast.Load()), ctx=ast.Load())
    ...: 
    ...: @macro
    ...: def bar(fun_def):
    ...:     return ast.FunctionDef(
    ...:         name=fun_def.name,
    ...:         body=fun_def.body,
    ...:         decorator_list=[],
    ...:         returns=fun_def.returns,
    ...:         args=ast.arguments(
    ...:             args=[ast.arg(arg='x', annotation=None)],
    ...:             vararg=None,
    ...:             kwonlyargs=[],
    ...:             kw_defaults=[],
    ...:             kwarg=None,
    ...:             defaults=[]))
    ...: 
    ...: @bar
    ...: def test():
    ...:     return x
    ...: 
    ...: x, y = 1, 2
    ...: print(foo(x + y))
    ...: print(test("HELLO WORLD, again"))
    ...: """)
got BinOp(left=Name(id='x', ctx=Load()), op=Add(), right=Name(id='y', ctx=Load()))

In [92]: exec(compile(a, "", "exec"))
HELLO WORLD -1
HELLO WORLD, again