Python AST 模块实战演示

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 模块的核心知识点整理成了下面的表格：

到此这篇关于Python AST 模块实战解析的文章就介绍到这了,更多相关python ast 模块内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

最新评论