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)
"""
原文:https://www.cnblogs.com/c-x-a/p/14832335.html