python ast NodeTransformer Transformer代码转换 代码插桩 新增node节点 编译生成新函数 类内函数更新

python ast NodeTransformer 代码插桩

import ast, astor, inspect

# 测试函数
def test_func(a,b):
    d = a + 2
    f = b + 5
    o = d + f
    return o

# 代码转换pass
class dump_trans(ast.NodeTransformer):
    def visit_Assign(self, node):
        dump_code = "print({})".format(node.targets[0].id)
        dump_node = ast.parse(dump_code).body[0]
        return node, dump_node

# 获取原函数ast
tast = ast.parse(inspect.getsource(test_func)).body[0]

# 进行代码插入
dt = dump_trans()
ntast = dt.visit(tast)

# ast 转换成 源代码 str
new_code = astor.to_source(ntast)
print(new_code)
'''
def test_func(a,b):
    d = a + 2
    print(d)
    f = b + 5
    print(f)
    o = d + f
    print(o)
    return o
'''

# 执行新代码 编译 生成 新函数对象
local={}
exec(new_code, {}, local)

print(local)

new_f = local["test_func"]

# 执行新函数
new_f(6, 9.0)
# 会有打印信息

装饰器内实现代码插桩 类内函数更新

import ast, astor, inspect

# 获取函数源代码
def get_func_code(func):
    lines = inspect.getsource(func).split('\n')
    col_lead = len(lines[0]) - len(lines[0].lstrip(" "))
    lines = [c[col_lead:] for c in  lines]
    code_str = "\n".join(lines)
    return code_str

# dump runtime 函数
def dump(name, var):
    print(name + " : ", var)

class inject_dump(ast.NodeTransformer):
    def __init__(self, global_val):
        self.global_val = global_val
        self.trans_func = {}

    def visit_Assign(self, node):
        node.value = self.visit(node.value)
        dump_code = "dump(\'{}\', {})".format(node.targets[0].id, node.targets[0].id)
        dump_node = ast.parse(dump_code).body[0]
        return node, dump_node
    
    def visit_Call(self, node):
        if isinstance(node.func, ast.Attribute) and isinstance(node.func.value, ast.Name) and\
            node.func.value.id == "self" and "self" in self.global_val:
            sub_func = getattr(self.global_val["self"], node.func.attr)
            ij_dump = inject_dump(self.global_val)
            fast = ast.parse(get_func_code(sub_func)).body[0]
            nfast = ij_dump.visit(fast)
            new_code = astor.to_source(nfast)
            print(new_code)
            local = {}
            exec(new_code, self.global_val, local)
            new_func = local[sub_func.__name__]
            # 更新类内的子函数  已经是非对象方法
            setattr(self.global_val["self"], sub_func.__name__, new_func)
            # 加 self参数
            return ast.Call(node.func, [ast.Name("self", node.func.ctx)] + node.args, node.keywords)
            
def transf(func):
    def _deco(*args, **kwargs):
        funn_args = inspect.signature(func)
        global_val={"dump":dump,}
        if "self" in funn_args.parameters:
            print(args[0])
            global_val["self"] = args[0]
            print(getattr(args[0],"s"))
        t = inject_dump(global_val)
        fast = ast.parse(get_func_code(func)).body[0]
        fast.decorator_list = [] # de decorate
        nfast = t.visit(fast)
        print(getattr(args[0],"s"))
        #print(ast.dump(nfast))
        #nfast = ast.fix_missing_locations(nfast)
        new_code = astor.to_source(nfast)
        print(new_code)
        local = {}
        exec(new_code, global_val, local)
        print(local)
        new_func = local[func.__name__]
        return new_func(*args, **kwargs)
    return _deco

@transf
def test_func(a:int, b:float):
    c = a + 1
    d = b + 1
    out = c + d
    return out

class tclass():
    def s(self, o):
        o = o -1
        return o
    @transf
    def test_func(self, a:int, b:float):
        c = a + 1
        d = b + 1
        out = c + d
        out = self.s(out)
        return out


test_func(3, 5.0)
tclass().test_func(3, 5.0)