AST - 像lisp一样自定义代码行为

前言

学common lisp(以下除非特殊需要说明的都简称lisp)以及用emacs的人都有一个体会 - lisp无所不能, 可以使用lisp修改lisp的行为. 什么意思呢?

我来举个例子. 我希望重置+的行为为实际意义的减法-. 看起来这是语言不可能完成的任务, 对lisp来说很简洁(我使用sbcl):

* (+ 1 1)    2  ; 正确结果  * (shadow '+)    T  * (defgeneric + (a &rest b))    #<STANDARD-GENERIC-FUNCTION + (0)>  * (defmethod + ((a number) &rest b) (apply 'cl:- a b))    #<STANDARD-METHOD + (NUMBER) {1002E43E73}>  * (+ 1 1)    0  ; 这里的加号的意义其实是我们所理解的`减号`

是不是很神奇?

那么对于python这种高级语言能不能做到呢? 答案是肯定的. 我们马上就来实现它

In [1]: import ast    In [2]: x = ast.parse('1 + 1', mode='eval')    In [3]: x.body.op = ast.Sub()    In [4]: eval(compile(x, '<string>', 'eval'))  Out[4]: 0

我想大家开始明白AST有多大能量了吧?

AST的故事

AST中文叫做抽象语法树, 也就是分析当前版本的python代码的语法, 用一种树的结构解析出来. 这个模块提供给我们一个在编译代码之前, 用python语言本身去修改.

它的作者是Armin Ronacher. 如果你听过或者觉得似曾相识, 对. 他就是mitsuhiko - flask的作者. 也是pocoo的leader之一(另外一个是看起来不知名的birkenfeld - 对我来说他很有名).

那么AST有什么意义呢? 但是有绝大多数人其实不了解也用不到这个模块, 为什么呢?

1. 出现需要对代码默认行为做更改的场景很少
2. 它主要用来做静态文件的检查, 比如pylint, pychecker，以及写flake8插件. 而我们平时的写代码都是在运行不需要进行预先的语法检查之类, 那么实际接触它就很难得了.

一些文章的索引

为了对本文有更深的理解可以看看以下文章

AST 模块：用 Python 修改 Python 代码这里对流程说的很好了. 可以直接读一下

模块代码也写得非常精炼, 可能不直接让你明白, 那么这时候可以看看

Abstract Syntax Trees, 这个时候我再强调一下作者吧, takluyver是ipython的核心开发成员, 他也参与了很多我们常用的开源项目, 比如pexpect和pandas

上面的2篇文章写了很多, 既有理解, 也有一些初级的用法.

我个人用它的例子

最近做的slack-alert. 先说它和AST的关系:

1. 我没有使用注册或者import的方式，而是直接去遍历文件, 找到符合我要求的函数当做一个任务需要执行的任务
2. 任务就要设置间隔, 那么会加某种格式的装饰器, 装饰器的参数就是间隔类型, 比如@deco(seconds=10)表示没十秒跑一次的意思
3. 我这样就可以放心的写plugin就好了, 我只关注任务本身的逻辑. 而这个装饰器(类似上面说的@deco), 它其实是不存在
4. 这个特殊格式的装饰器本身不存在没有关系, 因为我不会直接运行代码, 我只是把代码通过AST的处理, 解析出我要的任务和任务的执行间隔. 再去编译代码.

上代码:

class GetJobs(ast.NodeTransformer):        def __init__(self):  # 原来的ast.NodeTransformer其实没有__init__方法          self.jobs = []        def get_jobs(self):  # 一个方便的获得任务的方法          return self.jobs        def get_job_args(self, decorator):  # 这属于解析装饰器这个结构, 拿到执行的间隔          return {k.arg: k.value.n for k in decorator.keywords                  if k.arg in ('hours', 'seconds', 'minutes', 'days')                  and isinstance(k.value, ast.Num)}        def visit_FunctionDef(self, node):  # 这个visit_xxx的方法被重载的时候, 就会对这个类型的语法加一些特殊处理. 因为我设计的时候只有函数才有可能是任务          decorator_list = node.decorator_list  # 或者一个函数的装饰器列表          if not decorator_list:              return node  # 没有装饰器明显不是我想要的任务, 可能只是一个helper函数而已          decorator = decorator_list[0]  # 这里我把最外面的装饰器取出来看看是不是符合我要的格式          args = self.get_job_args(decorator)          if args:  # 当获得了适合的参数, 那么正确这个格式是正确的              node.decorator_list = decorator_list[1:] # 最外面的装饰器就是语法hack, 它不存在也没有意义，以后完成历史任务 去掉之              self.jobs.append((node.name, args))          return node      def find_jobs(path):      jobs = []      for root, dirs, files in os.walk(path):          for name in files:              file = os.path.join(root, name)              if not file.endswith('.py'):                  continue              with open(file) as f:                  expr_ast = ast.parse(f.read())  # 读文件, 解析                  transformer = GetJobs()                  sandbox = {}  # 其实就是把执行放在一个命名空间里面, 因为最后我还是会把任务编译执行的, 我在这里面存了执行后的环境                  exec(compile(transformer.visit(expr_ast),                               '<string>', 'exec'), sandbox)                  jobs.extend([(sandbox[j], kw) for j, kw in transformer.jobs])      return jobs

其实看起来不能完成的事情, 就是这么简单.