AST transformations with __future__-like module

In the previous article I wrote how-to add partial application with ... and piping with @ using AST transformations. However we needed to transform AST manually. For automatizing it I planned to use macropy but it doesn’t work with Python 3 and a bit too complicated. So I ended up with an idea to create __transformers__ module that work in a similar way with Python’s __future__ module. So code will look like:

from __transformers__ import ellipsis_partial, matmul_pipe

range(10) @ map(lambda x: x ** 2, ...) @ list @ print

So first of all for implementing it we need to extract enabled transformers names from code, it’s easy with ast.NodeVisitor, we just process all ImportForm nodes:

import ast


class NodeVisitor(ast.NodeVisitor):
    def __init__(self):
        self._found = []

    def visit_ImportFrom(self, node):
        if node.module == '__transformers__':
            self._found += [name.name for name in node.names]

    @classmethod
    def get_transformers(cls, tree):
        visitor = cls()
        visitor.visit(tree)
        return visitor._found

Let’s run it:

tree = ast.parse(code)

>>> print(NodeVisitor.get_transformers(tree))
['ellipsis_partial', 'matmul_pipe']

Next step is to define transformers. Transformer is just a Python module with transformer variable, that is instance of ast.NodeTransformer. For example transformer module for piping with matrix multiplication operator will be like:

import ast


class MatMulPipeTransformer(ast.NodeTransformer):
    def _replace_with_call(self, node):
        """Call right part of operation with left part as an argument."""
        return ast.Call(func=node.right, args=[node.left], keywords=[])

    def visit_BinOp(self, node):
        if isinstance(node.op, ast.MatMult):
            node = self._replace_with_call(node)
            node = ast.fix_missing_locations(node)

        return self.generic_visit(node)


transformer = MatMulPipeTransformer()

Now we can write function that extracts used transformers, imports and applies it to AST:

def transform(tree):
    transformers = NodeVisitor.get_transformers(tree)

    for module_name in transformers:
        module = import_module('__transformers__.{}'.format(module_name))
        tree = module.transformer.visit(tree)

    return tree

And use it on our code:

from astunparse import unparse

>>> unparse(transform(tree))
from __transformers__ import ellipsis_partial, matmul_pipe
print(list((lambda __ellipsis_partial_arg_0: map((lambda x: (x ** 2)), __ellipsis_partial_arg_0))(range(10)))

Next part is to automatically apply transformations on module import, for that we need to implement custom Finder and Loader. Finder is almost similar with PathFinder, we just need to replace Loader with ours in spec. And Loader is almost SourceFileLoader, but we need to run our transformations in source_to_code method:

from importlib.machinery import PathFinder, SourceFileLoader


class Finder(PathFinder):
    @classmethod
    def find_spec(cls, fullname, path=None, target=None):
        spec = super(Finder, cls).find_spec(fullname, path, target)
        if spec is None:
            return None

        spec.loader = Loader(spec.loader.name, spec.loader.path)
        return spec


class Loader(SourceFileLoader):
    def source_to_code(self, data, path, *, _optimize=-1):
        tree = ast.parse(data)
        tree = transform(tree)
        return compile(tree, path, 'exec',
                       dont_inherit=True, optimize=_optimize)

Then we need to put our finder in sys.meta_path:

import sys

def setup():
    sys.meta_path.insert(0, Finder)
    
setup()

And now we can just import modules that use transformers. But it requires some bootstrapping.

We can make it easier by creating __main__ module that will register module finder and run module or file:

from runpy import run_module
from pathlib import Path
import sys
from . import setup

setup()

del sys.argv[0]

if sys.argv[0] == '-m':
    del sys.argv[0]
    run_module(sys.argv[0])
else:
    # rnupy.run_path ignores meta_path for first import
    path = Path(sys.argv[0]).parent.as_posix()
    module_name = Path(sys.argv[0]).name[:-3]
    sys.path.insert(0, path)
    run_module(module_name)

So now we can run our module easily:

➜ python -m __transformers__ -m test   
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

➜ python -m __transformers__ test.py                 
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

And that’s all, you can try transformers by yourself with transformers package:

pip install transformers

Source code on github, package, previous part.