在Python中用AST包装一个带有装饰器的函数

import ast

# 定义一个装饰器函数
def decorator(func):
    def wrapper(*args, **kwargs):
        print("Before function call")
        result = func(*args, **kwargs)
        print("After function call")
        return result
    return wrapper

# 定义一个带有装饰器的函数
@decorator
def my_function():
    print("Inside my_function")

# 将带有装饰器的函数包装为AST
def wrap_function_with_decorator():
    # 解析源代码为抽象语法树
    code = """
@decorator
def my_function():
    print("Inside my_function")
"""
    tree = ast.parse(code)

    # 获取函数定义节点
    function_def = tree.body[0]

    # 创建装饰器节点
    decorator = ast.Name(id='decorator', ctx=ast.Load())
    decorator_call = ast.Call(func=decorator, args=[], keywords=[])

    # 将装饰器节点添加到函数定义节点的装饰器列表中
    function_def.decorator_list = [decorator_call]

    # 将修改后的抽象语法树转换为源代码
    modified_code = ast.unparse(tree)

    return modified_code

# 测试函数包装
print(wrap_function_with_decorator())