Modify a Python AST parsed from a method code node in order to rewrite certain specific calls

Question

I'm writing some tooling for online programming contexts.

Part of it is a test case checker which actually based on a set of pairs of (input, output) files are gonna check whether the solution method is actually working.

Basically, the solution method is expected to be defined as follow:

def solution(Nexter: inputs):
    # blahblah some code here and there
    n = inputs.next_int()
    sub_process(inputs)
    # simulating a print something
    yield str(n)

can be then translated (once the AST modifications) as:

def solution():
    # blahblah some code here and there
    n = int(input())
    sub_process()
    print(str(n))

Note: Nexter is a class defined to be whether a generator of user input() calls or carry out the expected inputs + some other goodies.

I'm aware of the issues related to converting back to source code from the AST (requires to rely on 3rd party stuff). I also know that there is a NodeTransformer class:

But its use remains unclear to me I don't know if I'm better off checking calls, expr, etc.

Here is below what I've ended up with:

signature = inspect.signature(iterative_greedy_solution)
if len(signature.parameters) == 1 and "inputs" in signature.parameters:
    parameter = signature.parameters["inputs"]
    annotation = parameter.annotation
    if Nexter == annotation:
        source = inspect.getsource(iterative_greedy_solution)
        tree = ast.parse(source)
        NexterInputsRewriter().generic_visit(tree)

class NexterInputsRewriter(ast.NodeTransformer):
    def visit(self, node):
        #???

This is definitely not the best design ever. Next time, I would probably go for the other way around (i.e. having a definition with simple user defined input() (and output, i.e. print(...)) and replacing them with test case inputs) when passing to a tester class asserting whether actual outputs are matching expecting ones.

To sum up this what I would like to achieve and I don't really know exactly how (apart of subclassing the NodeTransformer class):

Get rid of the solution function arguments
Modifiy the inputs calls in method body (as well as in the sub calls of methods also leveraging Nexter: inputs) in order to replace them with their actual user input() implementation, e.g. inputs.next_int() = int(input())

EDIT

Found that tool (https://python-ast-explorer.com/) that helps a lot to visualize what kind of ast.AST derivatives are used for a given function.

It is unclear to me what you really are trying to do. What I understand is that you want to start with the soure text for a python script to be tested, and you want to modify that script source to enable you to run automated tests against it. [Please confirm]. If I have it right, what you want for a script-modifying tool is a Program Transformation System. See my SO answer for how to modify a script at http://stackoverflow.com/a/39945762/120163 — Ira Baxter, Oct 24 '16 at 03:21
Actually it's more the other way around, my function definition is testable per se but I would like to convert it in order replace inputs: Nexter calls with their respective implementations and remove the inputs: Nexter as a parameter of the function definition. — Natalie Perret, Oct 24 '16 at 11:23
I interpret your answer as "I want to modify existing code to control how it gets inputs". If you aren't doing that to make it more testable, then why are you doing it? In any case, you want to modify the code in a regular way. My link suggestion stands. — Ira Baxter, Oct 24 '16 at 13:20
You'll find many examples of modifying the ast in: https://github.com/QQuick/Transcrypt/blob/master/transcrypt/modules/org/transcrypt/compiler.py . Search for 'on the fly'. Also you can run Transcrypt with the -dt option to generate a human readable syntax tree. This was how I found out how to do the transformations. It's a 3 step process: 1. Parse. 2. Generate modified node using info of the original one. 3. Visit it to generate modified code. — Jacques de Hooge, Jan 16 '17 at 11:43

score 0 · Answer 1 · answered Jul 31 '21 at 17:39

You can probably use NodeTransformer + ast.unparse() though it wouldn't be as effective as checking out some other 3rd party solutions considering it won't preserve any of your comments.

Here is an example transformation done by refactor (I'm the author), which is a wrapper layer around ast.unparse for doing easy source-to-source transformations through AST;

import ast

import refactor
from refactor import ReplacementAction

class ReplaceNexts(refactor.Rule):
    def match(self, node):
        # We need a call
        assert isinstance(node, ast.Call)

        # on an attribute (inputs.xxx)
        assert isinstance(node.func, ast.Attribute)

        # where the name for attribute is `inputs`
        assert isinstance(node.func.value, ast.Name)
        assert node.func.value.id == "inputs"

        target_func_name = node.func.attr.removeprefix("next_")

        # make a call to target_func_name (e.g int) with input()
        target_func = ast.Call(
            ast.Name(target_func_name),
            args=[
                ast.Call(ast.Name("input"), args=[], keywords=[]),
            ],
            keywords=[],
        )
        return ReplacementAction(node, target_func)


session = refactor.Session([ReplaceNexts])
source = """\
def solution(Nexter: inputs):
    # blahblah some code here and there
    n = inputs.next_int()
    sub_process(inputs)
    st = inputs.next_str()
    sub_process(st)
"""
print(session.run(source))

$ python t.py
def solution(Nexter: inputs):
    # blahblah some code here and there
    n = int(input())
    sub_process(inputs)
    st = str(input())
    sub_process(st)

Modify a Python AST parsed from a method code node in order to rewrite certain specific calls

EDIT

1 Answers1