修改从方法代码节点解析的 Python AST，以重写某些特定调用

Question

我正在为在线编程上下文编写一些工具。

它的一部分是一个测试用例检查器，它实际上基于一组（输入，输出）文件将检查解决方案方法是否实际工作。

基本上，求解方法预计定义如下：

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

然后可以翻译（一旦 AST 修改）为：

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

注意：Nexter 是一个定义为用户 input() 生成器是否调用或执行预期输入 + 其他一些好处的类。

我知道与从 AST 转换回源代码相关的问题（需要依赖第 3 方的东西）。我也知道有一个 NodeTransformer 类：

• http://greentreesnakes.readthedocs.io/en/latest/manipulating.html
• https://docs.python.org/3/library/ast.html#ast.NodeTransformer

但我仍然不清楚它的用途我不知道我是否最好检查调用、expr 等。

下面是我最终得到的结果：

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):
        #???

这绝对不是有史以来最好的设计。下一次，当传递给测试器类断言是否实际输出与预期匹配。

总结一下我想要实现的目标，但我真的不知道如何实现（除了子类化 NodeTransformer 类）：

• 去掉求解函数参数
• 修改方法主体中的输入调用（以及方法的子调用也利用 Nexter: 输入），以便将它们替换为实际的用户 input() 实现，例如inputs.next_int() = int(input())

编辑

找到了该工具 (https://python-ast-explorer.com/)，该工具对可视化用于给定函数的 ast.AST 导数有很大帮助。

Answer 1

您可能可以使用NodeTransformer + ast.unparse()，尽管它不会像查看其他一些第三方解决方案那样有效，因为它不会保留您的任何 cmets。

这是由refactor（我是作者）完成的示例转换，它是 ast.unparse 的包装层，用于通过 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)