移动测试开发 不要再手动批量替换了,使用 python AST 模块批量替换

opentest-oper@360.cn · 2022年07月13日 · 3418 次阅读

前言

在我们日常协作开发时,在团队内没有良好的规范或者 code review 机制时,经常会出现使用语义不明的变量,比如 a、b、c 等,使得代码可读性非常差,如果要将变量变更为具有语义的变量,批量替换容易替换错,而挨个手动替换也容易遗漏或者出错,因此,需要寻找快捷准确的方式处理这种情况。本文我们只针对 python 语言,首先我们先来了解一下 python 语言的编译过程。

我们整天和代码打交道,都知道高级语言分为解释型语言和编译型语言。解释型语言通过相应的解释器,将源代码翻译成目标代码(机器语言,也就是二进制形式),边解释边执行,因为边解释边执行的特点,因此执行效率比较低,且运行时不能脱离解释器;编译型语言通过编译器,将源代码编译成目标代码(机器语言),因此编译型语言可以脱离语言环境独立执行,使用方便且效率高,但每次更改代码都需要重新编译。本文我们详细介绍 Python(本文我们基于 CPython 解释器),python 和 Java 类似,解释器实际上分为两部分:编译器和虚拟机,先将代码编译成字节码,然后再由虚拟机执行。

在编译时,需要先经过语法分析,当代码出现语法错误时,就会在这个阶段抛出,接下来我们先了解一下语法分析的基础知识,编译的过程分为六个步骤,如下所示:

python 中可以使用 py_compile 模块将源代码编译成 PyCodeObject,PyCodeObject 进一步持久化到文件 pyc 中,解释器执行 pyc 文件,在 python 虚拟机中执行。

我们先来了解一下抽象语法树是什么?简单来说,抽象语法树经过语法分析(文法定义、文法分析以及消除左递归)后,选择某一条产生式进行展开后的结果。如果还不是很理解,后续我们会针对词法分析、语法分析、语义分析分别进行讲解。python 中可以使用 AST 模块,将代码转换为抽象语法树,接下来我们进入 python AST 模块的详解。

AST 基础知识

ast 模块官方链接: https://docs.python.org/3/library/ast.html#ast-helpers

python 官方对 ast 模块的解释如下:

The ast module helps Python applications to process trees of the Python abstract syntax grammar. The abstract syntax itself might change with each Python release; this module helps to find out programmatically what the current grammar looks like.

An abstract syntax tree can be generated by passing ast.PyCF_ONLY_AST as a flag to the compile() built-in function, or using the parse() helper provided in this module. The result will be a tree of objects whose classes all inherit from ast.AST. An abstract syntax tree can be compiled into a Python code object using the built-in compile() function.

更详细的关于节点的类型描述,可以参考官方文档,就不再赘述。

AST 实战

创建 AST 并优雅的输出
我们使用 astpretty 模块,将 ast 对象更优雅的输出,我们将文件 web_util.py 中的所有代码转换为 ast 对象,并优雅的输出,web_util.py 内容如下:

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

class opUtil:
    @staticmethod
    def add_two_num(a, b):
        return a + b

    @staticmethod
    def mul_two_num(a, b):
        print(a, b)
        return a * b

ast 模块可以将文件的 read 直接当做输入,解析并输出 ast 对象的代码如下:

# -*- encoding: utf-8 -*-
import ast
import astpretty

filename = "web_util.py"
f = open(filename)

ast_obj = ast.parse(f.read(), mode="exec")
astpretty.pprint(ast_obj)

输出结果为:

Module(
    body=[
        ClassDef(
            lineno=3,
            col_offset=0,
            name='opUtil',
            bases=[],
            keywords=[],
            body=[
                FunctionDef(
                    lineno=4,
                    col_offset=4,
                    name='add_two_num',
                    args=arguments(
                        args=[
                            arg(lineno=5, col_offset=20, arg='a', annotation=None),
                            arg(lineno=5, col_offset=23, arg='b', annotation=None),
                        ],
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=[],
                    ),
                    body=[
                        Return(
                            lineno=6,
                            col_offset=8,
                            value=BinOp(
                                lineno=6,
                                col_offset=15,
                                left=Name(lineno=6, col_offset=15, id='a', ctx=Load()),
                                op=Add(),
                                right=Name(lineno=6, col_offset=19, id='b', ctx=Load()),
                            ),
                        ),
                    ],
                    decorator_list=[Name(lineno=4, col_offset=5, id='staticmethod', ctx=Load())],
                    returns=None,
                ),
                FunctionDef(
                    lineno=8,
                    col_offset=4,
                    name='mul_two_num',
                    args=arguments(
                        args=[
                            arg(lineno=9, col_offset=20, arg='a', annotation=None),
                            arg(lineno=9, col_offset=23, arg='b', annotation=None),
                        ],
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=[],
                    ),
                    body=[
                        Expr(
                            lineno=10,
                            col_offset=8,
                            value=Call(
                                lineno=10,
                                col_offset=8,
                                func=Name(lineno=10, col_offset=8, id='print', ctx=Load()),
                                args=[
                                    Name(lineno=10, col_offset=14, id='a', ctx=Load()),
                                    Name(lineno=10, col_offset=17, id='b', ctx=Load()),
                                ],
                                keywords=[],
                            ),
                        ),
                        Return(
                            lineno=11,
                            col_offset=8,
                            value=BinOp(
                                lineno=11,
                                col_offset=15,
                                left=Name(lineno=11, col_offset=15, id='a', ctx=Load()),
                                op=Mult(),
                                right=Name(lineno=11, col_offset=19, id='b', ctx=Load()),
                            ),
                        ),
                    ],
                    decorator_list=[Name(lineno=8, col_offset=5, id='staticmethod', ctx=Load())],
                    returns=None,
                ),
            ],
            decorator_list=[],
        ),
    ],
)

遍历 AST 并修改节点

采用 ast.NodeTransformer 的方式遍历抽象语法树,我们在遍历的过程中,将函数中参数 a 命名改为更有意义的 first_num,这个场景在日常开发中很常见,经常会有变量名命名格式不规范,但手动改起来又容易遗漏或者改错,成本还是很高的。代码如下:

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

filename = "web_util.py"
f = open(filename)

ast_obj = ast.parse(f.read(), mode="exec")
astpretty.pprint(ast_obj)


class visitor_ast(ast.NodeTransformer):
    def generic_visit(self, node):
        print("ALL", type(node).__name__)
        fields = node._fields
        if "id" in fields and node.id == "a":
            print("field id", node.id)
            node.id = "first_num"
        ast.NodeVisitor.generic_visit(self, node)

    def visit_FunctionDef(self, node):
        ast.NodeVisitor.generic_visit(self, node)
        args_num = len(node.args.args)
        args = tuple([arg.arg for arg in node.args.args])
        func_log_stmt = ''.join(["print('calling func: %s', " % node.name, "'args:'", ", %s" * args_num % args, ')'])
        node.body.insert(0, ast.parse(func_log_stmt))

    def visit_arg(self, node):
        fields = node._fields
        if "arg" in fields and node.arg == "a":
            print("field arg", node.arg)
            node.arg = "first_num"
        print("ARG", type(node).__name__, node.arg)
        ast.NodeVisitor.generic_visit(self, node)


v = visitor_ast()
v.visit(ast_obj)
astpretty.pprint(ast_obj)
print(astunparse.unparse(ast_obj))

可以使用 astunparse 模块,将 ast 对象还原为代码,在这段代码中做了两件事:1、将参数 a 统一修改命名为 first_num;2、在函数中添加 print 日志。
在 ast 输出数据 15-18 行中,这几行表示方法的传入参数,在遍历节点时,可以通过 arg 取出入参。

在 ast 输出数据 65-80 行中,id 为 a 的,是参数 a 的引用。

因此,我们在遍历时,可以根据节点的 id 或 arg 挑出指定的参数,然后进行替换即可,最后把抽象语法树再转化为代码。遍历节点时,我们的代码 16-18 行中,将方法内对参数 a 的引用,都改为 first_num,在 30-32 行中,将函数的入参修改为 first_name,24-26 行中,在抽象语法树中添加打印日志的节点,第 40 行中,将抽象语法树转化为代码,参数 a 都被替换为 first_num,第 10 行新增了调用函数的日志输出,实现了我们想要的效果,最终转化的代码如下:

class opUtil():

    @staticmethod
    def add_two_num(first_num, b):
        print('calling func: add_two_num', 'args:', first_num, b)
        return (first_num + b)

    @staticmethod
    def mul_two_num(first_num, b):
        print('calling func: mul_two_num', 'args:', first_num, b)
        print(first_num, b)
        return (first_num * b)

总结

其实 AST 在我们日常的业务开发中极少用到,AST 模块作为代码辅助检查功能非常有意义,比如语法检查,调试错误等等,我们上面仅仅用来全局替换变量及打印日志,还可以缩小范围修改某个函数或者某个类里的。除此之外,还可以用于检测汉字、closure 检查等,后续我们有更多的使用案例,也会单独介绍。

如果觉得我的文章对您有用,请随意打赏。您的支持将鼓励我继续创作!
暂无回复。
需要 登录 后方可回复, 如果你还没有账号请点击这里 注册