PythonAst介绍及应用

PythonAst介绍及应⽤
Abstract Syntax Trees即抽象语法树。

Ast是python源码到字节码的⼀种中间产物，借助ast模块可以从语法树的⾓度分析源码结构。

此外，我们不仅可以修改和执⾏语法树，还可以将Source⽣成的语法树unparse成python源码。

因此ast给python源码检查、语法分析、修改代码以及代码调试等留下了⾜够的发挥空间。

1. AST简介
Python官⽅提供的CPython解释器对python源码的处理过程如下：
1. Parse source code into a parse tree (Parser/pgen.c)
2. Transform parse tree into an Abstract Syntax Tree (Python/ast.c)
3. Transform AST into a Control Flow Graph (Python/compile.c)
4. Emit bytecode based on the Control Flow Graph (Python/compile.c)
即实际python代码的处理过程如下：
源代码解析 --> 语法树 --> 抽象语法树(AST) --> 控制流程图 --> 字节码
上述过程在python2.5之后被应⽤。

python源码⾸先被解析成语法树，随后⼜转换成抽象语法树。

在抽象语法树中我们可以看到源码⽂件中的python的语法结构。

⼤部分时间编程可能都不需要⽤到抽象语法树，但是在特定的条件和需求的情况下，AST⼜有其特殊的⽅便性。

下⾯是⼀个抽象语法的简单实例。

Module(body=[
Print(
dest=None,
values=[BinOp( left=Num(n=1),op=Add(),right=Num(n=2))],
nl=True,
)])
2. 创建AST
2.1 Compile函数
先简单了解⼀下compile函数。

compile(source, filename, mode[, flags[, dont_inherit]])
source -- 字符串或者AST（Abstract Syntax Trees）对象。

⼀般可将整个py⽂件内容file.read()传⼊。

filename -- 代码⽂件名称，如果不是从⽂件读取代码则传递⼀些可辨认的值。

mode -- 指定编译代码的种类。

可以指定为 exec, eval, single。

flags -- 变量作⽤域，局部命名空间，如果被提供，可以是任何映射对象。

flags和dont_inherit是⽤来控制编译源码时的标志。

func_def = \
"""
def add(x, y):
return x + y
print add(3, 5)
"""
使⽤Compile编译并执⾏：
>>> cm = compile(func_def, '<string>', 'exec')
>>> exec cm
>>> 8
上⾯func_def经过compile编译得到字节码，cm即code对象，True == isinstance(cm, types.CodeType)。

compile(source, filename, mode, ast.PyCF_ONLY_AST) <==> ast.parse(source, filename='<unknown>', mode='exec')
2.2 ⽣成ast
使⽤上⾯的func_def⽣成ast.
r_node = ast.parse(func_def)
print astunparse.dump(r_node) # print ast.dump(r_node)
下⾯是func_def对应的ast结构：
Module(body=[
FunctionDef(
name='add',
args=arguments(
args=[Name(id='x',ctx=Param()),Name(id='y',ctx=Param())],
vararg=None,
kwarg=None,
defaults=[]),
body=[Return(value=BinOp(
left=Name(id='x',ctx=Load()),
op=Add(),
right=Name(id='y',ctx=Load())))],
decorator_list=[]),
Print(
dest=None,
values=[Call(
func=Name(id='add',ctx=Load()),
args=[Num(n=3),Num(n=5)],
keywords=[],
starargs=None,
kwargs=None)],
nl=True)
])
除了ast.dump，有很多dump ast的第三⽅库，如astunparse, codegen, unparse等。

这些第三⽅库不仅能够以更好的⽅式展⽰出ast结构，还能够将ast反向导出python source代码。

module Python version "$Revision$"
{
mod = Module(stmt* body)| Expression(expr body)
stmt = FunctionDef(identifier name, arguments args, stmt* body, expr* decorator_list)
| ClassDef(identifier name, expr* bases, stmt* body, expr* decorator_list)
| Return(expr? value)
| Print(expr? dest, expr* values, bool nl)| For(expr target, expr iter, stmt* body, stmt* orelse)
expr = BoolOp(boolop op, expr* values)
| BinOp(expr left, operator op, expr right)| Lambda(arguments args, expr body)| Dict(expr* keys, expr* values)| Num(object n) -- a number as a PyObject.
| Str(string s) -- need to specify raw, unicode, etc?| Name(identifier id, expr_context ctx)
| List(expr* elts, expr_context ctx)
-- col_offset is the byte offset in the utf8 string the parser uses
attributes (int lineno, int col_offset)
expr_context = Load | Store | Del | AugLoad | AugStore | Param
boolop = And | Or
operator = Add | Sub | Mult | Div | Mod | Pow | LShift | RShift | BitOr | BitXor | BitAnd | FloorDiv
arguments = (expr* args, identifier? vararg, identifier? kwarg, expr* defaults)
}
View Code
上⾯是部分摘⾃官⽹的 Abstract Grammar，实际遍历ast Node过程中根据Node的类型访问其属性。

3. 遍历AST
python提供了两种⽅式来遍历整个抽象语法树。

3.1 ast.NodeTransfer
将func_def中的add函数中的加法运算改为减法，同时为函数实现添加调⽤⽇志。

1class CodeVisitor(ast.NodeVisitor):
2def visit_BinOp(self, node):
3if isinstance(node.op, ast.Add):
4 node.op = ast.Sub()
5 self.generic_visit(node)
6
7def visit_FunctionDef(self, node):
8print'Function Name:%s'%
9 self.generic_visit(node)
10 func_log_stmt = ast.Print(
11 dest = None,
12 values = [ast.Str(s = 'calling func: %s' % , lineno = 0, col_offset = 0)],
13 nl = True,
14 lineno = 0,
15 col_offset = 0,
16 )
17 node.body.insert(0, func_log_stmt)
18
19 r_node = ast.parse(func_def)
20 visitor = CodeVisitor()
21 visitor.visit(r_node)
22# print astunparse.dump(r_node)
23print astunparse.unparse(r_node)
24exec compile(r_node, '<string>', 'exec')
运⾏结果：
Function Name:add
def add(x, y):
print'calling func: add'
return (x - y)
print add(3, 5)
calling func: add
-2
3.2 ast.NodeTransformer
使⽤NodeVisitor主要是通过修改语法树上节点的⽅式改变AST结构，NodeTransformer主要是替换ast中的节点。

既然func_def中定义的add已经被改成⼀个减函数了，那么我们就彻底⼀点，把函数名和参数以及被调⽤的函数都在ast中改掉，并且将添加的函数调⽤log写的更加复杂⼀些，争取改的⾯⽬全⾮：-）
1class CodeTransformer(ast.NodeTransformer):
2def visit_BinOp(self, node):
3if isinstance(node.op, ast.Add):
4 node.op = ast.Sub()
5 self.generic_visit(node)
6return node
7
8def visit_FunctionDef(self, node):
9 self.generic_visit(node)
10if == 'add':
11 = 'sub'
12 args_num = len(node.args.args)
13 args = tuple([arg.id for arg in node.args.args])
14 func_log_stmt = ''.join(["print 'calling func: %s', " % , "'args:'", ", %s" * args_num % args])
15 node.body.insert(0, ast.parse(func_log_stmt))
16return node
17
18def visit_Name(self, node):
19 replace = {'add': 'sub', 'x': 'a', 'y': 'b'}
20 re_id = replace.get(node.id, None)
21 node.id = re_id or node.id
22 self.generic_visit(node)
23return node
24
25 r_node = ast.parse(func_def)
26 transformer = CodeTransformer()
27 r_node = transformer.visit(r_node)
28# print astunparse.dump(r_node)
29 source = astunparse.unparse(r_node)
30print source
31# exec compile(r_node, '<string>', 'exec') # 新加⼊的node func_log_stmt 缺少lineno和col_offset属性
32exec compile(source, '<string>', 'exec')
33exec compile(ast.parse(source), '<string>', 'exec')
结果：
def sub(a, b):
print'calling func: sub', 'args:', a, b
return (a - b)
print sub(3, 5)
calling func: sub args: 3 5
-2
calling func: sub args: 3 5
-2
代码中能够清楚的看到两者的区别。

这⾥不再赘述。

4.AST应⽤
AST模块实际编程中很少⽤到，但是作为⼀种源代码辅助检查⼿段是⾮常有意义的；语法检查，调试错误，特殊字段检测等。

上⾯通过为函数添加调⽤⽇志的信息是⼀种调试python源代码的⼀种⽅式，不过实际中我们是通过parse整个python⽂件的⽅式遍历修改源码。

4.1 汉字检测
下⾯是中⽇韩字符的unicode编码范围
CJK Unified Ideographs
Range: 4E00— 9FFF
Number of characters: 20992
Languages: chinese, japanese, korean, vietnamese
使⽤ unicode 范围 \u4e00 - \u9fff 来判别汉字，注意这个范围并不包含中⽂字符(e.g. u'；' == u'\uff1b') .
下⾯是⼀个判断字符串中是否包含中⽂字符的⼀个类CNCheckHelper：
1class CNCheckHelper(object):
2# 待检测⽂本可能的编码⽅式列表
3 VALID_ENCODING = ('utf-8', 'gbk')
4
5def _get_unicode_imp(self, value, idx = 0):
6if idx < len(self.VALID_ENCODING):
7try:
8return value.decode(self.VALID_ENCODING[idx])
9except:
10return self._get_unicode_imp(value, idx + 1)
11
12def _get_unicode(self, from_str):
13if isinstance(from_str, unicode):
14return None
15return self._get_unicode_imp(from_str)
16
17def is_any_chinese(self, check_str, is_strict = True):
18 unicode_str = self._get_unicode(check_str)
19if unicode_str:
20 c_func = any if is_strict else all
21return c_func(u'\u4e00' <= char <= u'\u9fff'for char in unicode_str)
22return False
接⼝is_any_chinese有两种判断模式，严格检测只要包含中⽂字符串就可以检查出，⾮严格必须全部包含中⽂。

下⾯我们利⽤ast来遍历源⽂件的抽象语法树，并检测其中字符串是否包含中⽂字符。

1class CodeCheck(ast.NodeVisitor):
2
3def__init__(self):
4 _checker = CNCheckHelper()
5
6def visit_Str(self, node):
7 self.generic_visit(node)
8# if node.s and any(u'\u4e00' <= char <= u'\u9fff' for char in node.s.decode('utf-8')):
9if _checker.is_any_chinese(node.s, True):
10print'line no: %d, column offset: %d, CN_Str: %s' % (node.lineno, node.col_offset, node.s)
11
12 project_dir = './your_project/script'
13for root, dirs, files in os.walk(project_dir):
14print root, dirs, files
15 py_files = filter(lambda file: file.endswith('.py'), files)
16 checker = CodeCheck()
17for file in py_files:
18 file_path = os.path.join(root, file)
19print'Checking: %s' % file_path
20 with open(file_path, 'r') as f:
21 root_node = ast.parse(f.read())
22 checker.visit(root_node)
上⾯这个例⼦⽐较的简单，但⼤概就是这个意思。

关于CPython解释器执⾏源码的过程可以参考官⽹描述：
4.2 Closure 检查
⼀个函数中定义的函数或者lambda中引⽤了⽗函数中的local variable，并且当做返回值返回。

特定场景下闭包是⾮常有⽤的，但是也很容易被误⽤。

关于python闭包的概念可以参考我的另⼀篇⽂章：
这⾥简单介绍⼀下如何借助ast来检测lambda中闭包的引⽤。

代码如下：
1class LambdaCheck(ast.NodeVisitor):
2
3def__init__(self):
4 self.illegal_args_list = []
5 self._cur_file = None
6 self._cur_lambda_args = []
7
8def set_cur_file(self, cur_file):
9assert os.path.isfile(cur_file)， cur_file
10 self._cur_file = os.path.realpath(cur_file)
11
12def visit_Lambda(self, node):
13"""
14 lambda 闭包检查原则：
15只需检测lambda expr body中args是否引⽤了lambda args list之外的参数
16"""
17 self._cur_lambda_args =[a.id for a in node.args.args]
18print astunparse.unparse(node)
19# print astunparse.dump(node)
20 self.get_lambda_body_args(node.body)
21 self.generic_visit(node)
22
23def record_args(self, name_node):
24if isinstance(name_node, ) and name_node.id not in self._cur_lambda_args:
25 self.illegal_args_list.append((self._cur_file, 'line no:%s' % name_node.lineno, 'var:%s' % name_node.id))
26
27def _is_args(self, node):
28if isinstance(node, ):
29 self.record_args(node)
30return True
31if isinstance(node, ast.Call):
32 map(self.record_args, node.args)
33return True
34return False
35
36def get_lambda_body_args(self, node):
37if self._is_args(node): return
38# for cnode in ast.walk(node):
39for cnode in ast.iter_child_nodes(node):
40if not self._is_args(cnode):
41 self.get_lambda_body_args(cnode)
遍历⼯程⽂件：
1 project_dir = './your project/script'
2for root, dirs, files in os.walk(project_dir):
3 py_files = filter(lambda file: file.endswith('.py'), files)
4 checker = LambdaCheck()
5for file in py_files:
6 file_path = os.path.join(root, file)
7 checker.set_cur_file(file_path)
8 with open(file_path, 'r') as f:
9 root_node = ast.parse(f.read())
10 checker.visit(root_node)
11 res = '\n'.join([' ## '.join(info) for info in checker.illegal_args_list])
12print res
View Code
由于Lambda(arguments args, expr body)中的body expression可能⾮常复杂，上⾯的例⼦中仅仅处理了⽐较简单的body expr。

可根据⾃⼰⼯程特点修改和扩展检查规则。

为了更加⼀般化可以单独写⼀个visitor类来遍历lambda节点。

Ast的应⽤不仅限于上⾯的例⼦，限于篇幅，先介绍到这⾥。

期待ast能帮助你解决⼀些⽐较棘⼿的问题。