Python AST 模块实战演示_Python

python 的 ast（abstract syntax tree，抽象语法树）模块是一个用于处理 python 代码的强大工具。它允许你将代码转换为一种结构化的树形表示，从而可以进行分析、修改甚至生成新的代码。这对于理解代码的内部结构、构建开发工具或进行代码自动化处理非常有帮助。

🔍 什么是抽象语法树（ast）

在深入了解 ast 模块之前，我们先简单理解一下抽象语法树（ast）的概念。当 python 解释器执行代码时，它首先需要理解代码的结构。这个过程大致是：源代码 -> 词法分析（生成令牌流） -> 语法分析（生成 ast） -> 字节码 -> 执行。ast 就是源代码抽象语法结构的树状表示，它过滤掉了像空格、注释这类非本质的细节，专注于代码的逻辑结构。每个节点代表代码中的一个结构（例如，一个表达式、一个语句、一个函数定义等）。

ast 模块的核心作用就是在这棵“语法树”上工作，让你能在代码被编译成字节码之前，洞察和操作其本质。

🛠️ ast 模块的核心用法

ast 模块提供了一系列函数和类来创建、遍历和修改 ast。

1. 解析代码生成 ast

使用 ast.parse() 函数可以将字符串形式的 python 代码解析成一棵 ast 的根节点（通常是 ast.module 节点）。

import ast
code = """
def greet(name):
    print(f"hello, {name}!")
"""
tree = ast.parse(code)  # 得到 ast

2. 查看 ast 结构

生成 ast 后，可以使用 ast.dump() 函数将其以文本形式打印出来，以便查看整个树的结构。

print(ast.dump(tree, indent=4))

这会输出一个结构化的文本，展示所有的节点、它们的属性以及嵌套关系。

3. 遍历 ast

要分析 ast，你需要遍历它的节点。ast 模块提供了两种主要方式：

使用 ast.nodevisitor 类：这是最常用和推荐的方法。你可以创建一个继承自 ast.nodevisitor 的类，并为感兴趣的节点类型定义 visit_ 方法（例如 visit_functiondef 用于访问函数定义节点）。在方法中，你可以通过 self.generic_visit(node) 来继续遍历当前节点的子节点。

class myvisitor(ast.nodevisitor):
    def visit_functiondef(self, node):
        print(f"found function: {node.name}")
        self.generic_visit(node)  # 继续遍历子节点
visitor = myvisitor()
visitor.visit(tree)

使用 ast.walk() 函数：这个函数会递归地遍历 ast 中的所有节点，但不关心节点的层级关系。它返回一个生成器，适合当你需要找到所有特定类型的节点时使用。

for node in ast.walk(tree):
    if isinstance(node, ast.functiondef):
        print(node.name)

4. 修改 ast

除了分析，你还可以修改 ast。这需要通过继承 ast.nodetransformer 类来实现。它的使用方式和 nodevisitor 类似，但关键区别在于，visit_ 方法需要返回一个节点。你可以返回：

原节点：不做任何修改。
新节点：替换原节点。
none：删除该节点。

一个经典的例子是将代码中的所有加法操作 + 替换为减法操作 -：

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)

重要提示：在创建新节点替换旧节点后，如果新节点缺少源代码位置信息（如行号、列偏移量），需要使用 ast.fix_missing_locations() 函数来修复，否则在编译时可能会出错。

5. 将 ast 编译回代码

修改完 ast 后，你可以通过内置的 compile() 函数将其编译成可执行的 python 字节码，然后使用 exec() 或 eval() 来运行它。

# 将修改后的 ast 编译成代码对象
code_obj = compile(new_tree, filename='<string>', mode='exec')
# 执行代码对象
exec(code_obj)

此外，从 python 3.9 开始，标准库提供了 ast.unparse() 函数，可以直接将 ast 节点转换回可读的 python 代码字符串。对于更早的版本，可以使用第三方库如 astor。

6. 安全地求值表达式

ast 模块还提供了一个非常实用的函数 ast.literal_eval()。它可以安全地计算一个包含 python 字面量（如字符串、数字、元组、列表、字典、布尔值等）的表达式字符串，并返回结果。与内置的 eval() 不同，它不会执行任意代码，因此安全得多，非常适合处理来自不可信来源的数据。

safe_result = ast.literal_eval("[1, 2, 3]")  # 结果是列表 [1, 2, 3]

💻 大厂笔试面试如何考察 ast

在大厂的笔试或面试中，对 ast 模块的考察通常不会要求你死记硬背所有的节点类型，而是更注重实践应用能力、对 python 机制的理解深度以及解决实际问题的思路。

常见的考察方向

基础概念理解
- 可能会问：ast 在 python 解释过程中处于哪个阶段？它和字节码有什么区别？
- 或者：ast.literal_eval() 和 eval() 有什么区别？为什么前者更安全？
代码分析与审计
- 静态代码检查：让你编写一个简单的检查器，使用 ast.nodevisitor 来遍历代码，找出潜在问题。例如，检测是否使用了不安全的函数（如 eval、exec），或者检查代码风格（如变量命名规范）。
- 代码复杂度分析：通过分析函数定义、循环、条件分支等节点的数量和嵌套关系，来估算代码的复杂度。
代码转换与自动化
- 这是考察的重点和难点。面试官可能会给出一个具体的代码重构任务，让你使用 ast.nodetransformer 来实现自动化。例如：
重命名：将代码中所有特定变量名或函数名进行批量替换。
逻辑替换：如前面提到的，将加法操作改为减法。
api 升级：假设某个库的 api 发生了变化，需要你写一个工具自动将旧代码中的函数调用更新为新形式。
与网络安全结合
- 尤其是在安全相关的岗位面试中。可能会给出一段代码，要求你分析其中可能存在的安全漏洞（如 sql 注入、命令注入点），这需要你能够通过 ast 分析代码的数据流和控制流。

实例分析：面试题模拟

题目：请你使用 ast 模块，编写一个简单的静态分析工具，用于检测一段 python 代码中是否使用了 eval() 函数。如果使用了，则输出警告信息。

考察点：

是否掌握 ast 的基本解析和遍历操作。
是否熟悉 ast.nodevisitor 的使用。
能否识别函数调用节点（ast.call）并判断其函数名。

参考实现：

import ast
code = """
x = 1
result = eval('1 + 1')
print(eval('2+2'))
"""
class evaldetector(ast.nodevisitor):
    def visit_call(self, node):
        # 检查节点是否是一个函数调用，并且函数名是一个标识符（name）且id为'eval'
        if isinstance(node.func, ast.name) and node.func.id == 'eval':
            print(f"warning: potential use of eval() found at line {node.lineno}")
        # 继续遍历子节点，以查找嵌套调用等情况
        self.generic_visit(node)
tree = ast.parse(code)
detector = evaldetector()
detector.visit(tree)

输出：

warning: potential use of eval() found at line 4
warning: potential use of eval() found at line 5

这个例子演示了如何使用 ast.nodevisitor 来访问代码中的函数调用节点（ast.call），并根据条件（函数名为 eval）进行判断和输出。

💎 总结与核心知识点

为了帮助你更好地记忆，我将 ast 模块的核心知识点整理成了下面的表格：

核心概念/操作	关键函数/类	说明与用途
解析代码	`ast.parse(source)`	将源代码字符串解析为 ast 根节点（`ast.module`）。
查看结构	`ast.dump(node)`	将 ast 节点以字符串形式输出，用于调试。
遍历 ast	`ast.nodevisitor`	通过继承此类并定义 `visit_xxx` 方法来有选择地访问节点。
遍历 ast	`ast.walk(node)`	递归遍历 ast 中的所有节点，不保留层级信息。
修改 ast	`ast.nodetransformer`	通过继承此类，在 `visit_xxx` 方法中返回新节点来修改 ast。
编译执行	`compile(tree, ...)`	将 ast 编译为可执行的代码对象。
反解析	`ast.unparse(node)` (python 3.9+)	将 ast 节点转换回等价的 python 代码字符串。
安全求值	`ast.literal_eval()`	安全地求值字面量表达式（字符串、列表、数字等），避免任意代码执行风险。
修复位置	`ast.fix_missing_locations(node)`	为新建或修改的节点补充行号等位置信息，确保能正确编译。