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__importellipsis_partial,matmul_piperange(10)@map(lambdax:x**2,...)@list@printSo 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:
importastclassNodeVisitor(ast.NodeVisitor):def__init__(self):self._found=[]defvisit_ImportFrom(self,node):ifnode.module=='__transformers__':self._found+=[name.namefornameinnode.names]@classmethoddefget_transformers(cls,tree):visitor=cls()visitor.visit(tree)returnvisitor._foundLet’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:
importastclassMatMulPipeTransformer(ast.NodeTransformer):def_replace_with_call(self,node):"""Call right part of operation with left part as an argument."""returnast.Call(func=node.right,args=[node.left],keywords=[])defvisit_BinOp(self,node):ifisinstance(node.op,ast.MatMult):node=self._replace_with_call(node)node=ast.fix_missing_locations(node)returnself.generic_visit(node)transformer=MatMulPipeTransformer()Now we can write function that extracts used transformers, imports and applies it to AST:
deftransform(tree):transformers=NodeVisitor.get_transformers(tree)formodule_nameintransformers:module=import_module('__transformers__.{}'.format(module_name))tree=module.transformer.visit(tree)returntreeAnd use it on our code:
fromastunparseimportunparse>>>unparse(transform(tree))from__transformers__importellipsis_partial,matmul_pipeprint(list((lambda__ellipsis_partial_arg_0:map((lambdax:(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:
fromimportlib.machineryimportPathFinder,SourceFileLoaderclassFinder(PathFinder):@classmethoddeffind_spec(cls,fullname,path=None,target=None):spec=super(Finder,cls).find_spec(fullname,path,target)ifspecisNone:returnNonespec.loader=Loader(spec.loader.name,spec.loader.path)returnspecclassLoader(SourceFileLoader):defsource_to_code(self,data,path,*,_optimize=-1):tree=ast.parse(data)tree=transform(tree)returncompile(tree,path,'exec',dont_inherit=True,optimize=_optimize)Then we need to put our finder in sys.meta_path:
importsysdefsetup():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:
fromrunpyimportrun_modulefrompathlibimportPathimportsysfrom.importsetupsetup()delsys.argv[0]ifsys.argv[0]=='-m':delsys.argv[0]run_module(sys.argv[0])else:# rnupy.run_path ignores meta_path for first importpath=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