[PYTHON] Example of code rewriting by ast.NodeTransformer

It's not a hassle to write a Converter I will write a little trial content. As a reminder, I rarely touch ast in Python.

What I wanted to do

For example, I want to generate a function that doubles the argument v and returns it when the string " v * 2 " is given. (It is known that the name of the argument is v)

Realization policy

For Single Statement

As in the above example, if the given string is a Single Statement, it's very easy, just use a string operation to make it lambda and eval.

single.py

text = "v * 2"
func = eval("lambda v:" + text)

Support for Multiple Statements

If the given string can contain more than one expression, the above method cannot be used. For example

from System.Windows import Thickness
Thickness(v, 0, 0, 0)

Consider the case where is specified. For the time being, the only thing that comes to my mind is to use a string operation to create a function def and then execute it. (This code only returns None for the generated function, but I'll talk about that later)

multi.py

text = """
from System.Windows import Thickness
Thickness(v, 0, 0, 0)
"""
import re
source = "def test(v):\n" + re.sub("^", "  ", text, flags=re.M)
exec(source)
return test

This is fine in most cases, but for example, when a multi-line literal is given, there is a problem that the contents of the literal are indented by the simple indent processing by the above regular expression. ..

"""value of v is:
{v}""".format(v=v)

It's a rare case that you don't really need to consider, but it's a bit uncomfortable to know that there is a bad pattern, so as an alternative, we will make it a function def by ast conversion.

1. First, generate an ast for an empty function def called def test (v): pass.
2. Replace the Pass node part of the generated ast with the ast node generated from the given string.

The procedure.

It looks like this when written in code.

functionize.py

import ast

class _MakeFunction(ast.NodeTransformer):
    def __init__(self, body):
        self.body = body

    def run(self):
        template_ast = ast.parse('def test(v): pass', mode='exec')
        return self.visit(template_ast)
 
    def visit_Pass(self, node):
        return ast.parse(self.body, mode='exec').body
 
if __name__ == '__main__':
    body = "from System.Windows import Thickness\nThickness(v, 0, 0, 0)"
    functionized_ast = _MakeFunction(body).run()
    exec(compile(functionized_ast, '<string>', mode='exec'))

Return value generation

Now it's a function, but as I wrote earlier, there is no return value as it is. So, using ast conversion, rewrite it so that the result of the last evaluated expression is returned as a return value like Ruby. The procedure is as follows.

1. Add the code _retval_ = None to the beginning of function
2. Add the code return _retval_ to the end of function
3. Rewrite all the expressions that appear in function to the assignment statement for _retval_.

Let's write a Transformer that does this in code.

functionize.py

class _ReturnLastResult(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        self.generic_visit(node)
        retval_assign = \
                ast.Assign(targets=[ast.Name(id='_retval_', ctx=ast.Store())],
                           value=ast.NameConstant(value=None))               
        retval_return = \
                ast.Return(value=ast.Name(id='_retval_', ctx=ast.Load()))
        node.body.insert(0, retval_assign)
        node.body.append(retval_return)
        return ast.fix_missing_locations(node)

    def visit_Expr(self, node):
        target = ast.Name(id='_retval_', ctx=ast.Store())
        assign = ast.copy_location(
                    ast.Assign(targets=[target], value=node.value), node)
        return assign

code

The final code looks like this. https://gist.github.com/wonderful-panda/8a22b74248a60cc8bb22

After all, I decided to consider only Single Statement in the original entry and did not use this.