AST transformations with __future__-like module

Inthe 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 .