时间戳

首页 > 技术文章 > 使用ast(抽象语法树)在代码中植入埋点

使用ast(抽象语法树)在代码中植入埋点

oaks 2020-07-18 22:57 原文

功能说明

在代码执行过程中收集一些日志，但是这个操作是与业务无关的，需要根据运行环境来决定
是否要执行这些操作，
一个解决办法：

在代码中加点儿收集日志的标记，比如“注释”，既不会对改变原有代码的结构又能实现功能，
入侵性比较低。

那就需要解释器能识别“注释”并翻译成代码执行了，比如对于以下代码：

# -*- encoding: utf-8 -*-

metrics = {
    "f1": 0.9,
    "acc": 0.8
}
'''@myclient.send_metrics(metrics)'''

希望解释器执行时能把

'''@myclient.send_metrics(metrics)'''

当成代码：

myclient.send_metrics(metrics)

也就是整体代码要翻译成：

# -*- encoding: utf-8 -*-
import myclient

metrics = {
    "f1": 0.9,
    "acc": 0.8
}

myclient.send_metrics(metrics)

简单粗暴的方法替换字符串就可以了，还可以用ast（Abstract Syntax Trees）模块处理这个问题。

代码块有各种类型：

  • 方法定义
  • 变量申明、赋值
  • 计算表达式
  • ...

在python中，语句：

'''@myclient.send_metrics(metrics)'''

表示定义了一个字符串对象，有时候这种形式用于注释，但不是真的注释，解释器不会把它忽略，
只要从语法树中找到这样的代码节点，把它的内容从字符串申明变成方法调用就可以了。

实现

1. 找到埋点在语法树中的代码节点

# -*- encoding: utf-8 -*-
import ast, astunparse

code = \
"""
# -*- encoding: utf-8 -*-

metrics = {
    "f1": 0.9,
    "acc": 0.8
}
'''@myclient.send_metrics(metrics)'''
"""

ast_tree = ast.parse(code)  # 解析成语法树

for node in ast_tree.body:
    if isinstance(node, ast.Expr) and isinstance(node.value, ast.Str):  # 字符串定义表达式
        anno_content = node.value.s
        if anno_content.startswith('@myclient.send_metrics'):  # 发现标记
            anno_code = anno_content[1:]  # 去掉@
            print("这个node=%s中带有埋点，埋点代码:\n%s" % (str(node), anno_code))

执行结果：

这个node=<_ast.Expr object at 0x120336250>中带有埋点，埋点代码:
myclient.send_metrics(metrics)

2. 把埋点的字符串变成方法调用并生成新代码

# -*- encoding: utf-8 -*-
import ast, astunparse, astor

code = \
"""
# -*- encoding: utf-8 -*-

metrics = {
    "f1": 0.9,
    "acc": 0.8
}
'''@myclient.send_metrics(metrics)'''
"""

ast_tree = ast.parse(code)  # 解析成语法树

for node in ast_tree.body:
    if isinstance(node, ast.Expr) and isinstance(node.value, ast.Str):   # 字符串定义表达式
        anno_content = node.value.s
        if anno_content.startswith('@myclient.send_metrics'):  # 发现标记
            anno_code = anno_content[1:]  # 去掉@
            print("这个node=%s中带有埋点，埋点代码:\n%s" % (str(node), anno_code))
            anno_ast_tree = ast.parse(anno_code)  # 用"注释"中的代码生成一个新的语法树
            anno_expr = anno_ast_tree.body[0]  # 获取语法树中的表达式，也就是 myclient.send_metrics(metrics)
            anno_call = anno_expr.value  # 获取表达式中的方法调用
            node.value = anno_call  # 把赋值语句变成方法调用语句

# 在第1行导入一下  myclient
import_nni = ast.Import(names=[ast.alias(name='myclient', asname=None)])
nodes = ast_tree.body
nodes.insert(0, import_nni)

print("最终代码: ")
print(astor.to_source(ast_tree))

输出是：

这个node=<_ast.Expr object at 0x1214156d0>中带有埋点，埋点代码:
myclient.send_metrics(metrics)
最终代码: 
import myclient
metrics = {'f1': 0.9, 'acc': 0.8}
myclient.send_metrics(metrics)

3. 执行生成的代码

python 提供的compile、exec、eval 等函数都可以执行。

4. 上述代码的问题：

  1. 只遍历了根节点下的那一层节点，也就是不支持在方法中使用这种特殊注释，这个需要递归遍历所有语法节点
  2. import 是直接加入第一行的，可以遍历代码查看是否有导入，决定添加import语句以及其位置

推荐阅读