Python 重构神器：告别 ast 模块的痛苦，三行代码搞定复杂批量修改

1. 引言：Python 开发者的“缩进噩梦”

想象一下，你接手了一个古老的 Python 项目，Tech Lead 让你把所有的 unittest 断言风格： assert foo == bar 全部迁移成： self.assertEqual(foo, bar)

你打开 IDE，准备写一个正则表达式。突然你意识到：

缩进地狱：Python 是缩进敏感的，正则很难处理多层嵌套里的缩进。
多行问题：如果 assert 的内容太长换行了怎么办？
上下文：你只想改 class Test(unittest.TestCase) 里的代码，不想误伤普通的 assert 语句。

以前，你可能需要花半天时间去研究 Python 自带的 ast 库，写一个 100 行的 NodeTransformer。现在，用 ast-grep，你只需要 30 秒。

2. 概念拆解：它懂 Python 的语法

生活类比：Word 查找替换 vs. 智能改卷助手

正则表达式 就像 Word 里的“查找替换”。它只认识字母，根本不在乎你是在函数里，还是在注释里，更不懂 Python 的缩进规则。

ast-grep 就像一位懂 Python 语法的助教。你告诉它：“把所有测试用例里的‘相等断言’改一下写法”。它能看懂哪里是函数体，哪里是装饰器，哪里是换行续行。

核心优势：

无视缩进：你写的 Pattern 不需要关心缩进，ast-grep 会自动适配目标代码的缩进层级。
智能匹配：它能区分字符串里的 assert 和真正的关键字 assert。

3. 动手实战：从 `print` 到 `logging`

我们先看一个经典场景：把调试用的 print 换成生产环境的 logger.info。

第一步：安装

ast-grep 既提供了命令行工具，也提供了 Python 绑定（可以直接在 Python 脚本里调用）。

Bash

# 安装命令行工具 (推荐作为入口)
pip install ast-grep-cli

# 或者安装 Python 绑定库 (用于写复杂脚本)
pip install ast-grep

第二步：实战演练 (CLI 方式)

创建一个 legacy.py：

Python

def process_data(data):
    # 这里的 print 应该被替换
    print(f"Processing {data}")
    
    if not data:
        # 即使有缩进，也能完美处理
        print("Error: No data")
        return

在终端运行：

Bash

sg run --lang python \
  --pattern 'print($MSG)' \
  --rewrite 'logger.info($MSG)' \
  legacy.py

见证奇迹： legacy.py 瞬间变为：

Python

def process_data(data):
    # 这里的 print 应该被替换
    logger.info(f"Processing {data}")
    
    if not data:
        # 即使有缩进，也能完美处理
        logger.info("Error: No data")
        return

注意：它不仅替换了文本，还完美保留了第二行 logger 的缩进！

4. 进阶深潜：用 Python 脚本操作 Python 代码

CLI 适合简单替换，但如果你需要复杂的逻辑（比如：只替换函数名为 test_ 开头的函数里的断言），这时候 Python 绑定 (pip install ast-grep) 就派上用场了。

场景：迁移 Unittest 断言

我们需要把 assert a == b 变成 self.assertEqual(a, b)，但前提是这两个变量必须是简单的变量，不能是复杂的函数调用。

创建一个脚本 refactor.py：

Python

from ast_grep_py import SgRoot

# 模拟一段源代码
code = """
class MyTest(unittest.TestCase):
    def test_basic(self):
        x = 1
        y = 1
        assert x == y  # 我们想改这个
        assert complex_call() == 2 # 我们不想改这个
"""

# 1. 解析代码
sg = SgRoot(code, "python")
root = sg.root()

# 2. 查找模式：assert $A == $B
# pattern 语法就是 Python 语法加上 $ 通配符
matches = root.find_all(pattern="assert $A == $B")

for match in matches:
    # 3. 获取捕获的变量
    var_a = match.get_match("A")
    var_b = match.get_match("B")
    
    # 4. 加上自定义逻辑：只替换纯粹的标识符(变量名)，跳过复杂表达式
    # kind 检查是 ast-grep 的强项
    if var_a.kind() == "identifier" and var_b.kind() == "identifier":
        # 5. 执行替换
        new_text = f"self.assertEqual({var_a.text()}, {var_b.text()})"
        replace_range = match.range()
        # 这里只是演示逻辑，真实场景可以使用 match.replace(new_text)
        print(f"Found match! Replacing line {replace_range.start.line + 1}:")
        print(f"  Old: {match.text()}")
        print(f"  New: {new_text}")

运行结果：它只会匹配到 assert x == y，而聪明地跳过了 assert complex_call() == 2，因为我们在 Python 脚本里加了 kind() == "identifier" 的判断逻辑。

这比写正则安全了一万倍！

5. Python 专属技巧与陷阱

技巧：匹配装饰器 (Decorators)

Python 的装饰器写法很灵活，用 ast-grep 怎么匹配？

Pattern:

Python

@$DECO
def $FUNC():
  $$$

$DECO 匹配装饰器名（如 pytest.fixture）。
$FUNC 匹配函数名。
$$$ (三个美元符号) 是多行通配符，匹配函数体内的所有代码。

陷阱：字典的 Key

在 Python 中，{ key: val } 和 { "key": val } 在 AST 层面是不同的。

如果 pattern 写 { a: 1 }，它匹配不到 { "a": 1 }。
解决方法：如果你想同时匹配，可能需要写两条规则，或者使用 YAML 配置中的 any 组合规则。

6. 总结与延伸

一句话总结：对于 Python 开发者，ast-grep 是连接“简单正则”和“复杂 AST 模块”的桥梁，它让你用 Python 的原生语法去重构 Python 代码，同时解决了缩进敏感的千古难题。

posted @ 2026-01-08 08:57 Swizard 阅读(5) 评论(0) 收藏举报

刷新页面返回顶部