ast_utils.py 文件源码

def apply_ast_transform(func, ast_transformer, *,
                        keep_original=True, globals_dict=None, debug=0):
    """
    Apply the AST transform class to a function

    Args:
        keep_original: True to retain the old function in attribute `.f_original`
        globals_dict: pass any external function in your NodeTransformer into this
    """
    if (inspect.isclass(ast_transformer)
        and issubclass(ast_transformer, DecoratorAST)):
        ast_transformer = ast_transformer()
    else:
        assert isinstance(ast_transformer, DecoratorAST)

    old_ast = get_func_ast(func)
    # _, starting_line = inspect.getsourcelines(func)
    if debug:
        print("======= OLD AST =======")
        ast_print(old_ast)
    visitor = ast_transformer
    new_ast = visitor.visit(old_ast)
    if debug:
        print("======= NEW AST =======")
        ast_print(new_ast)
    ast.fix_missing_locations(new_ast)
    co = compile(new_ast, '<ast_demo>', 'exec')
    fake_locals = {}
    # exec will define the new function into fake_locals scope
    # this is to avoid conflict with vars in the real locals()
    # https://stackoverflow.com/questions/24733831/using-a-function-defined-in-an-execed-string-in-python-3
    exec(co, globals_dict, fake_locals)
    new_f = fake_locals[func.__name__]
    new_f.f_original = func if keep_original else None
    return new_f