Python 抽象语法树（AST）操作：代码分析与自动重构

好的，各位观众，欢迎来到“Python AST 操作：代码分析与自动重构” 讲座现场！今天，咱们要聊聊 Python 代码背后的秘密，以及如何用 AST 这个神器来玩转你的代码。

开场白：代码界的“透视眼”

想象一下，你拥有了一双能看穿 Python 代码表象，直达其本质的“透视眼”。这双眼睛能让你理解代码的真正含义，甚至可以自动修改和优化代码。是不是很酷？

这双“透视眼”就是抽象语法树 (Abstract Syntax Tree, AST)。

代码分析和自动重构是软件开发中至关重要的环节。代码分析帮助我们理解代码的结构、发现潜在的错误和安全漏洞，而自动重构则可以自动化地改进代码质量、提高代码的可维护性。AST 在这两方面都扮演着核心角色。

第一部分：什么是 AST？代码的骨架

AST 是一种代码的树状表示形式。它忽略了代码中的一些细节，比如空格、注释等等，只关注代码的结构和语义。简单来说，AST 就像是代码的骨架，它揭示了代码的组织方式和逻辑关系。

咱们举个例子：

x = 1 + 2 * 3

这行简单的代码，对应的 AST 长什么样呢？

用 ast 模块打印出来看看：

import ast

code = "x = 1 + 2 * 3"
tree = ast.parse(code)
print(ast.dump(tree, indent=4))

输出：

Module(
    body=[
        Assign(
            targets=[
                Name(id='x', ctx=Store())
            ],
            value=BinOp(
                left=Constant(value=1),
                op=Add(),
                right=BinOp(
                    left=Constant(value=2),
                    op=Mult(),
                    right=Constant(value=3)
                )
            ),
            type_comment=None
        )
    ],
    type_ignores=[]
)

别被这一堆东西吓到！咱们来解读一下：

Module：表示整个模块（也就是整个 Python 文件）。
Assign：表示赋值语句。
- targets：赋值的目标，这里是变量 x。
- value：赋的值，这里是一个二元运算表达式。
BinOp：表示二元运算。
- left：左边的操作数。
- op：运算符，Add 表示加法，Mult 表示乘法。
- right：右边的操作数。
Constant：表示常量，这里是数字 1、2 和 3。
Name：表示变量名，这里是 x。

可以看到，AST 将代码分解成一个个节点，并用树状结构表示它们之间的关系。是不是有点像剥洋葱，一层一层地揭开代码的秘密？

第二部分：AST 的基本操作：增删改查

有了 AST，我们就可以对代码进行各种操作了。主要有以下几种：

遍历 (Traversal)：访问 AST 中的每个节点，这是进行代码分析的基础。
查找 (Searching)：在 AST 中查找特定的节点，比如查找所有函数定义、查找所有变量等等。
修改 (Modification)：修改 AST 中的节点，比如修改变量名、修改运算符等等。
生成 (Generation)：根据 AST 生成新的代码。

咱们一个个来看：

1. 遍历 (Traversal)

遍历 AST 通常使用 ast.NodeVisitor 类。咱们定义一个简单的 visitor，打印出所有遇到的节点类型：

import ast

class MyVisitor(ast.NodeVisitor):
    def visit(self, node):
        print(type(node))
        self.generic_visit(node) # 继续遍历子节点

code = "x = 1 + 2 * 3"
tree = ast.parse(code)
visitor = MyVisitor()
visitor.visit(tree)

输出：

<class 'ast.Module'>
<class 'ast.Assign'>
<class 'ast.Name'>
<class 'ast.BinOp'>
<class 'ast.Constant'>
<class 'ast.BinOp'>
<class 'ast.Constant'>
<class 'ast.Mult'>
<class 'ast.Constant'>
<class 'ast.Add'>

NodeVisitor 提供了一个 visit 方法，它会被 AST 中的每个节点调用。generic_visit 方法会继续遍历节点的子节点。

当然，我们也可以针对特定类型的节点进行处理。比如，只打印出所有的变量名：

import ast

class VariableVisitor(ast.NodeVisitor):
    def visit_Name(self, node):
        print(f"变量名: {node.id}")

code = "x = 1 + 2 * y"
tree = ast.parse(code)
visitor = VariableVisitor()
visitor.visit(tree)

输出：

变量名: x
变量名: y

注意这里我们定义了 visit_Name 方法，它只会被 Name 类型的节点调用。

2. 查找 (Searching)

查找可以基于遍历来实现，但更方便的方式是使用 ast.walk 函数。它可以迭代 AST 中的所有节点。

比如，查找所有函数定义：

import ast

code = """
def add(x, y):
    return x + y

def subtract(x, y):
    return x - y
"""

tree = ast.parse(code)

for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef):
        print(f"函数名: {node.name}")

输出：

函数名: add
函数名: subtract

3. 修改 (Modification)

修改 AST 通常使用 ast.NodeTransformer 类。它允许你修改 AST 中的节点，并返回修改后的节点。

比如，将所有的加法运算替换成减法运算：

import ast

class AddToSubTransformer(ast.NodeTransformer):
    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub() # 将加法替换成减法
        return self.generic_visit(node) # 继续遍历子节点

code = "x = 1 + 2 * 3 + 4"
tree = ast.parse(code)
transformer = AddToSubTransformer()
new_tree = transformer.visit(tree)
print(ast.dump(new_tree, indent=4))

输出：

Module(
    body=[
        Assign(
            targets=[
                Name(id='x', ctx=Store())
            ],
            value=BinOp(
                left=BinOp(
                    left=Constant(value=1),
                    op=Sub(),
                    right=BinOp(
                        left=Constant(value=2),
                        op=Mult(),
                        right=Constant(value=3)
                    )
                ),
                op=Sub(),
                right=Constant(value=4)
            ),
            type_comment=None
        )
    ],
    type_ignores=[]
)

可以看到，所有的 Add 节点都被替换成了 Sub 节点。

4. 生成 (Generation)

修改完 AST 之后，我们需要将它转换回代码。这可以使用 ast.unparse 函数（Python 3.9+）或者 astor 库。

import ast
import astor # 需要安装：pip install astor

# 假设我们已经有了 new_tree (修改后的 AST)
code = "x = 1 + 2 * 3 + 4"
tree = ast.parse(code)

class AddToSubTransformer(ast.NodeTransformer):
    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub() # 将加法替换成减法
        return self.generic_visit(node) # 继续遍历子节点

transformer = AddToSubTransformer()
new_tree = transformer.visit(tree)

new_code = astor.to_source(new_tree)
print(new_code)

输出：

x = 1 - 2 * 3 - 4

可以看到，代码中的加法运算都被替换成了减法运算。

第三部分：代码分析的应用：静态类型检查

AST 在代码分析方面有很多应用，比如静态类型检查、代码风格检查、安全漏洞检测等等。咱们以静态类型检查为例，来演示一下 AST 的威力。

Python 是一种动态类型语言，这意味着变量的类型是在运行时确定的。静态类型检查则是在编译时（或者在运行之前）检查代码中的类型错误。

咱们定义一个简单的类型检查器：

import ast

class TypeChecker(ast.NodeVisitor):
    def __init__(self):
        self.context = {} # 存储变量的类型信息

    def visit_Assign(self, node):
        if isinstance(node.targets[0], ast.Name):
            var_name = node.targets[0].id
            if isinstance(node.value, ast.Constant):
                if isinstance(node.value.value, int):
                    self.context[var_name] = "int"
                elif isinstance(node.value.value, str):
                    self.context[var_name] = "str"
            else:
                self.generic_visit(node)

    def visit_BinOp(self, node):
        self.generic_visit(node)
        if isinstance(node.left, ast.Name) and isinstance(node.right, ast.Name):
            left_type = self.context.get(node.left.id)
            right_type = self.context.get(node.right.id)
            if left_type and right_type and left_type != right_type:
                print(f"类型错误: {node.left.id} ({left_type}) 和 {node.right.id} ({right_type}) 类型不匹配")

这个类型检查器做了以下事情：

visit_Assign：检查赋值语句，如果右边是常量，则将变量的类型存储到 context 中。
visit_BinOp：检查二元运算，如果左右两边的变量类型不匹配，则打印错误信息。

咱们来测试一下：

code = """
x = 1
y = "hello"
z = x + y
"""

tree = ast.parse(code)
checker = TypeChecker()
checker.visit(tree)

输出：

类型错误: x (int) 和 y (str) 类型不匹配

可以看到，类型检查器成功地发现了 x 和 y 类型不匹配的错误。

第四部分：自动重构的应用：代码清理

AST 在自动重构方面也有很多应用，比如代码格式化、代码简化、代码优化等等。咱们以代码清理为例，来演示一下 AST 的应用。

代码清理的目标是去除代码中不必要的冗余，提高代码的可读性。比如，去除未使用的变量、去除重复的代码等等。

咱们定义一个简单的代码清理器，去除未使用的变量：

import ast

class UnusedVariableRemover(ast.NodeTransformer):
    def __init__(self):
        self.used_vars = set() # 存储已使用的变量

    def visit_Name(self, node):
        if isinstance(node.ctx, ast.Load):
            self.used_vars.add(node.id)
        return node

    def visit_Assign(self, node):
        if isinstance(node.targets[0], ast.Name):
            var_name = node.targets[0].id
            if var_name not in self.used_vars:
                return None # 删除未使用的变量
        return self.generic_visit(node)

这个代码清理器做了以下事情：

visit_Name：记录所有已使用的变量。
visit_Assign：如果赋值语句的目标变量没有被使用，则删除该赋值语句。

咱们来测试一下：

import ast
import astor

code = """
x = 1
y = 2
z = x + 3
print(z)
"""

tree = ast.parse(code)
remover = UnusedVariableRemover()
new_tree = remover.visit(tree)
new_code = astor.to_source(new_tree)
print(new_code)

输出：

x = 1
z = x + 3
print(z)

可以看到，未使用的变量 y 及其赋值语句被成功地删除了。

总结：AST，代码世界的瑞士军刀

AST 是一个非常强大的工具，它可以让你深入理解代码的本质，并对代码进行各种操作。无论是代码分析还是自动重构，AST 都是不可或缺的利器。

优点： 精确、可靠、可定制。
缺点： 学习曲线陡峭，需要一定的编程基础。

表格总结：AST 相关模块和库

模块/库	功能
`ast`	Python 内置的 AST 模块，提供 AST 的基本操作
`astor`	将 AST 转换成代码
`gast`	兼容不同 Python 版本的 AST
`pyflakes`	代码风格检查
`pylint`	代码质量检查
`flake8`	代码风格和错误检查

结束语：让代码更上一层楼

希望今天的讲座能让你对 Python AST 有更深入的了解。掌握 AST，你就能像一位代码界的魔法师一样，随心所欲地操控你的代码，让你的代码更上一层楼！

感谢大家的观看！咱们下期再见！

发表回复 取消回复

发表回复取消回复