我的 python 程序无法重写 AST 中的节点

My python program fails to rewrite a node in AST

我正在尝试用 raise ValueError 替换 print。但是这个程序 returns 只有 TypeError: expected some sort of expr, but got <ast.Raise object at 0x000001DF779469A0>。如何修复以下代码以实现我的目标?

import ast


class MyVisitor(ast.NodeTransformer):

    def visit_Name(self, node: ast.Name):
        if node.id == 'print':
            result = ast.parse("raise ValueError").body[0]
            result.lineno = node.lineno
            result.col_offset = node.col_offset
            return result
        return node


program = """
f = 4 + 4
print("sth")
"""
tree = ast.parse(program)
my_visitor = MyVisitor()
tree = my_visitor.visit(tree)
code = compile(ast.fix_missing_locations(tree), "filename", "exec")
exec(code)

改为访问Expr节点,可以替换打印语句 使用 raise valueError 语句:

import ast

class MyVisitor(ast.NodeTransformer):
    def visit_Expr(self, node: ast.Expr):
        print(f"{node.__class__ =}")
        if node.__class__ != ast.Expr:
            return node
        if node.value.__class__ != ast.Call:
            return node
        if node.value.func.id == 'print':
            result = ast.parse("raise ValueError").body[0]
            result.lineno = node.lineno
            result.col_offset = node.col_offset
            return result
        return node

program = """
f = 4 + 4
print("sth")
"""
tree = ast.parse(program)
atv_out = ast.dump(tree)
my_visitor = MyVisitor()
tree = my_visitor.visit(tree)
code = compile(ast.fix_missing_locations(tree), "filename", "exec")
exec(code)
# raises ValueError

如果您使用 python3.9+,您实际上可以使用 refactor,它的界面比 NodeTransformer 简单得多。例如,您的问题的解决方案可以是这样的;

import ast
import refactor

class ReplacePrint(refactor.Rule):
    def match(self, node):
        assert isinstance(node, ast.Call)
        assert isinstance(node.func, ast.Name)
        assert node.func.id == 'print'
        
        value_error = ast.Raise(exc=ast.Name('ValueError', ast.Load()))
        return refactor.ReplacementAction(node, value_error)

program = """
f = 4 + 4
print("sth")
"""

session = refactor.Session(rules=[ReplacePrint])
print(session.run(program))