python ast

import ast

import astor

# 初始代码

source = """
index=0
def some_function(param):
    if param == 0:
       return case_0(param)
    elif param < 0:
       return negative_case(param)
    for i in range(5):
        print(i)   
    return all_other_cases(param)
"""


class UseAst(astor.TreeWalk):
    def pre_body_name(self):
        body = self.cur_node
        for i, child in enumerate(body[:]):
            self.__name = None
            # 继续遍历当前的node
            self.walk(child)
            if self.__name is not None:
                #添加日志输出
                logger_statement = ast.Expr(ast.Call(func=(ast.Attribute(value=ast.Name(id='logger'), attr='info')),

                                                    args=[ast.Str("Calling {}".format(self.__name), ctx=ast.Load())],
                                                    keywords=[]
                                                    ))
                body.insert(i, logger_statement)
        self.__name = None
        return True

    def pre_Call(self):
        # 获取调用函数的函数名
        if isinstance(self.cur_node.func, ast.Name):
            self.__name = self.cur_node.func.id
        return True

    def pre_For(self):
        # 所有的For循环的节点都走这
        node = self.cur_node
        body = node.body
        parent = self.parent
        add_statement = ast.parse('items = []')
        parent.insert(1, add_statement)
        add_statement = ast.parse('items.append(i)')
        body.insert(0, add_statement)


tree = ast.parse(source)

walker = UseAst()
walker.walk(tree)
body = tree.body
#加入导包
body.insert(0, ast.ImportFrom(module='loguru', names=[ast.alias(name="logger", asname=None)], level=0))
print(astor.to_source(tree))

"""
最终代码

from loguru import logger
index = 0


def some_function(param):
    if param == 0:
        logger.info('Calling case_0')
        return case_0(param)
    elif param < 0:
        logger.info('Calling negative_case')
        return negative_case(param)
    items = []
    logger.info('Calling all_other_cases')
    for i in range(5):
        items.append(i)
        print i
    return all_other_cases(param)
"""