Python AST 抽象语法树

2021-01-10  本文已影响0人  默写年华Antifragile

Abstract Sytax Tree
暂时用到的原因:在模型量化中,需要量化某些操作符带来的运算效果,比如 '+', '-','*', '/' 等等,这些就需要对源代码进行查询,因此就要需要将python解释器已经将源代码转化为运行的类后,再翻转回源代码
参考:
https://docs.python.org/3/library/ast.html#ast.NodeTransformer
https://www.cnblogs.com/yssjun/p/10069199.html
Abstract Syntax Trees即抽象语法树。Ast是python源码到字节码的一种中间产物,借助ast模块可以从语法树的角度分析源码结构。此外,我们不仅可以修改和执行语法树,还可以将Source生成的语法树unparse成python源码。因此ast给python源码检查、语法分析、修改代码以及代码调试等留下了足够的发挥空间。

1. AST简介

Python官方提供的CPython解释器对python源码的处理过程如下:

即实际python代码的处理过程如下:

上述过程在python2.5之后被应用。python源码首先被解析成语法树,随后又转换成抽象语法树。在抽象语法树中我们可以看到源码文件中的python的语法结构。
大部分时间编程可能都不需要用到抽象语法树,但是在特定的条件和需求的情况下,AST又有其特殊的方便性。

下面是一个抽象语法的简单实例。

func_def = \
""" 
def add(x, y):
    return x+y
    
print(add(3,5))
"""
print(func_def)

其中 三引号可以根据书写的方式智能换行,输出如下:


def add(x, y):
    return x+y
    
print(add(3,5))

image.png

2. 创建AST

2.1 compile(source, filename, mode[, flags[, dont_inherit]])

这是python自带的函数

>>> cm = compile(func_def, filename='<string>', mode='exec')
>>> exec(cm)
8
>>> type(cm)
code

上面func_def经过compile编译得到字节码,cm即code对象,True == isinstance(cm, types.CodeType)。

2.2 生成AST

>>> cm1 = ast.parse(func_def,filename='<unknown>', mode='exec')
>>> type(cm1)
_ast.Module
>>> ast.dump(cm1)
(
body=[
    FunctionDef(name='add', 
                args=arguments(
                                args=[arg(arg='x', annotation=None), arg(arg='y', annotation=None)], 
                                vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]
                              ), 
                body=[Return(
                                value=BinOp(left=Name(id='x', ctx=Load()), op=Add(), right=Name(id='y', ctx=Load()))
                            )
                     ], 
                decorator_list=[], 
                returns=None), 
    Expr(value=Call(
                    func=Name(id='print', ctx=Load()), 
                    args=[Call(func=Name(id='add', ctx=Load()), args=[Num(n=3), Num(n=5)], keywords=[])], 
                    keywords=[])
                    )
    ]
)

可以看到,这里对源代码进行了解析

3. 遍历语法树

python提供了两种方式来遍历整个语法树

generic_visit(node)
This visitor calls visit() on all children of the node. Note that child nodes of nodes that have a custom visitor method won’t be visited unless the visitor calls generic_visit() or visits them itself.

3.1 ast.NodeVisitor

比如 我们将 func_def 的 add 函数中的加法运算改为减法

class CodeVisitor(ast.NodeVisitor):
    def visit_BinOp(self, node):# 这个函数的访问是由于 Visit_FunctionDef的先访问再generic_visit才访问的
        print('Bin')            # 如果Visit_FunctionDef中没有generic_visit的话,则这个函数是不会访问的
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub()
            
        self.generic_visit(node)
    
    def visit_FunctionDef(self, node):
        print('Function Name: %s'% node.name)
        self.generic_visit(node) # FunctionDef中还包含有 BinOp,因此会进去visit BinOP
        
    def visit_Call(self, node):
        print("call") 
        self.generic_visit(node) # 因为AST的Call中还包含有一个Call,因此会重复再访问一次 
    
        
r_node = ast.parse(func_def)
visitor = CodeVisitor()
visitor.visit(r_node) # 这里的visit函数会根据 node 的语法树去遍历里面的函数,

输出:

Function Name: add
Bin
call
call

3.2 ast.NodeTransformer

A NodeVisitor subclass that walks the abstract syntax tree and allows modification of nodes

使用NodeVisitor主要是通过修改语法树上节点的方式改变AST结构,NodeTransformer主要是替换ast中的节点。

class CodeTransformer(ast.NodeTransformer):
    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub()
        self.generic_visit(node)
        return node

    def visit_FunctionDef(self, node):
        self.generic_visit(node) # 这里表示先去访问里面的children node        
        if node.name == 'add':
            node.name = 'sub'
        args_num = len(node.args.args)
        
        args_num = len(node.args.args)
        args = tuple([arg.arg for arg in node.args.args])
        print(str(args))
        func_log_stmt = ''.join(["print('calling func: %s', " % node.name, "'args:'", ", %s" * args_num % args ,')'])
        node.body.insert(0, ast.parse(func_log_stmt))
        
        #func_log_stmt = ''.join(["print 'calling func: %s', " % node.name, "'args:'", ", %s" * args_num % args])
        #node.body.insert(0, ast.parse(func_log_stmt))

        return node

    def visit_Name(self, node):
        replace = {'add': 'sub', 'x': 'a', 'y': 'b'}
        re_id = replace.get(node.id, None)
        node.id = re_id or node.id
        self.generic_visit(node)
        return node
    
    def visit_arg(self, node):
        self.generic_visit(node)
        replace = {'x':'a', 'y':'b'}
        node.arg = replace[node.arg]
        return node
        
        
r_node = ast.parse(func_def)
transformer = CodeTransformer()
r_node = transformer.visit(r_node)
#print(astunparse.dump(r_node))
source = astunparse.unparse(r_node) # astunparse 一般python不自带,需要conda 或者 pip安装
print(source)

输出:

('a', 'b')


def sub(a, b):
    print('calling func: sub', 'args:', a, b)
    return (a - b)
print(sub(3, 5))

可以看加入了一个print语句,然后将变量名字由 x, y 改为了 a, b


Keep in mind that if the node you’re operating on has child nodes you must either transform the child nodes yourself or call the generic_visit() method for the node first.


Don’t use the NodeVisitor if you want to apply changes to nodes during traversal. For this a special visitor exists (NodeTransformer) that allows modifications.

上一篇 下一篇

猜你喜欢

热点阅读