转载

使用 ast 模块魔改 Python 代码

上一篇博文「魔改 pip 的快乐与痛苦」中，我提到魔改开源代码，一个需要时刻考虑的事情，就是如何方便地合并上游变更。为了合并上游变更，我们要尽量克制自己乱写代码的冲动，不要作出那些使得合并代码变得困难重重的事情，比如像我曾经干出的修改包目录的蠢事。

为了修正我犯下的错，我决定在 mpip 代码仓库中保持基本 pip 一样的包目录，改换包目录这种事情放到编译过程里做。于是我不仅要在编译过程里把代码复制一份，还得把包目录从 pip 移动到 mpip，然后把代码里各种和 pip 相关的包目录、文件路径、常量定义、环境变量，总之一切有可能让 pip 和 mpip 之间不清不楚互相干扰的，都得改干净，这样才能让 mpip 和 pip 和平共处。

一开始我的想法是用 sed 直接做文本替换，虽然也作出了一个能用的版本，但是总是感觉慎得慌，不知道有没有该改的没改，不该改的改了，以后改代码会不会让现在的替换脚本失效之类的。总之就是解决方案太低级，确定性和可扩展性都不上台面。

于是我决定不使用简单的文本替换，而是直接在语法树上作文章，遍历 pip 代码中的语法元素，把变量名、字符串、函数等等的名字给改了，这样子基本上就能同时满足我想要的确定性和可扩展性了。

Python 提供了一个叫 ast 的模块，能够将 Python 代码 parse 成 AST，还提供了一个 NodeTransformer 让开发者自定义遍历语法树的行为。

于是实现这个替换工具就很简单了，把要替换的语法元素的 visit 方法实现，然后针对不同语法元素的不同属性作字符串替换，替换好了再用 astunparse 这个模块把 AST 变回代码，写回到文件里。

感觉有点 Lisp 的意思，用 Python 修改 Python。

原文 https://blog.jamespan.me/posts/modify-python-code-with-ast

正文到此结束