将 AST 节点转换为向量/数字

Question

您好，我有一组抽象语法树 (AST) 中的 Python 3 代码。我已经尝试了好几天，以找出将节点转换为可用向量/数字表示的最佳方法。

例如，这里是转储的一个 AST（没有注释字段）：

Module([Import([alias('time', None), alias('sys', None), alias('pygame', None)]), Import([alias('random', None)]), Import([alias('sequence', None)]), Assign([Name('S_SIZE', Store()), Tuple([Name('S_WID', Store()), Name('S_HGT', Store())], Store())], Tuple([Num(600), Num(400)], Load())), Assign([Name('screen', Store())], Call(Attribute(Attribute(Name('pygame', Load()), 'display', Load()), 'set_mode', Load()), [Name('S_SIZE', Load())], [])), Assign([Name('NUMB_COUNT', Store())], Num(200)), Assign([Name('nlist', Store())], ListComp(BinOp(Call(Attribute(Name('random', Load()), 'random', Load()), [], []), Mult(), Name('S_HGT', Load())), [comprehension(Name('_', Store()), Call(Name('range', Load()), [Name('NUMB_COUNT', Load())], []), [], 0)])), Assign([Name('num', Store())], Call(Attribute(Name('sequence', Load()), 'NumGroup', Load()), [Name('nlist', Load()), Name('S_WID', Load()), Name('S_HGT', Load())], [])), FunctionDef('draw_all', arguments([], None, [], [], None, []), [Expr(Call(Attribute(Name('screen', Load()), 'fill', Load()), [Tuple([Num(0), Num(0), Num(0)], Load())], [])), Expr(Call(Attribute(Name('num', Load()), 'draw', Load()), [Name('screen', Load())], [])), Expr(Call(Attribute(Attribute(Name('pygame', Load()), 'display', Load()), 'flip', Load()), [], []))], [], None), FunctionDef('bubble_sort', arguments([arg('nlist', None), arg('i', None), arg('end_ind', None)], None, [], [], None, []), [If(Compare(Name('i', Load()), [Eq()], [Name('end_ind', Load())]), [Return(Tuple([Num(0), BinOp(Name('end_ind', Load()), Sub(), Num(1))], Load()))], []), If(Compare(Attribute(Subscript(Name('nlist', Load()), Index(Name('i', Load())), Load()), 'val', Load()), [Gt()], [Attribute(Subscript(Name('nlist', Load()), Index(BinOp(Name('i', Load()), Add(), Num(1))), Load()), 'val', Load())]), [Expr(Call(Attribute(Name('nlist', Load()), 'swap', Load()), [Name('i', Load()), BinOp(Name('i', Load()), Add(), Num(1))], []))], []), Return(Tuple([BinOp(Name('i', Load()), Add(), Num(1)), Name('end_ind', Load())], Load()))], [], None), If(Compare(Name('__name__', Load()), [Eq()], [Str('__main__')]), [Expr(Call(Attribute(Name('pygame', Load()), 'init', Load()), [], [])), Assign([Name('i', Store())], Num(0)), Assign([Name('end_ind', Store())], BinOp(Call(Name('len', Load()), [Name('num', Load())], []), Sub(), Num(1))), While(Num(1), [For(Name('event', Store()), Call(Attribute(Attribute(Name('pygame', Load()), 'event', Load()), 'get', Load()), [], []), [If(Compare(Attribute(Name('event', Load()), 'type', Load()), [Eq()], [Attribute(Name('pygame', Load()), 'QUIT', Load())]), [Expr(Call(Attribute(Name('sys', Load()), 'exit', Load()), [], []))], [])], []), For(Name('n', Store()), Name('num', Load()), [Expr(Call(Attribute(Name('n', Load()), 'set_color', Load()), [Tuple([Num(255), Num(255), Num(255)], Load())], []))], []), If(Compare(Name('end_ind', Load()), [NotEq()], [Num(0)]), [Expr(Call(Attribute(Subscript(Name('num', Load()), Index(Name('i', Load())), Load()), 'set_color', Load()), [Tuple([Num(0), Num(255), Num(0)], Load())], []))], []), If(Compare(Name('end_ind', Load()), [NotEq()], [Num(0)]), [Assign([Tuple([Name('i', Store()), Name('end_ind', Store())], Store())], Call(Name('bubble_sort', Load()), [Name('num', Load()), Name('i', Load()), Name('end_ind', Load())], []))], []), Expr(Call(Attribute(Name('num', Load()), 'update', Load()), [], [])), Expr(Call(Name('draw_all', Load()), [], []))], [])], [])])

我想把它做成这样的东西，这样我就可以把它输入 TensorFlow：

[1,2,3,1,2,13,56,12,53,41,31...etc]

我找到了所有节点的副本（转换成字典）：

NODE_LIST = [
'Module','Interactive','Expression','FunctionDef','ClassDef','Return',
'Delete','Assign','AugAssign','Print','For','While','If','With','Raise',
'TryExcept','TryFinally','Assert','Import','ImportFrom','Exec','Global',
'Expr','Pass','Break','Continue','attributes','BoolOp','BinOp','UnaryOp',
'Lambda','IfExp','Dict','Set','ListComp','SetComp','DictComp',
'GeneratorExp','Yield','Compare','Call','Repr','Num','Str','Attribute',
'Subscript','Name','List','Tuple','Load','Store','Del',
'AugLoad','AugStore','Param','Ellipsis','Slice','ExtSlice','Index','And','Or',
'Add','Sub','Mult','Div','Mod','Pow','LShift','RShift','BitOr','BitXor',
'BitAnd','FloorDiv','Invert','Not','UAdd','USub','Eq','NotEq','Lt',
'LtE','Gt','GtE','Is','IsNot','In','NotIn','comprehension','ExceptHandler',
'arguments','keyword','alias']

NODE_MAP = {x: i for (i, x) in enumerate(NODE_LIST)}

例如，

{'Module':1,'Interactive':2,...etc}

我尝试了 ASTWalkers 和生成器，但我仍然找不到实现此目的的好方法。任何帮助表示赞赏:)

编辑： 我想我可能正在寻找ast.NodeVisitor的def visit_Name

class ToInteger(ast.NodeVisitor):

    def visit_Name(self, node):
        print(node.id)
        print(NODE_MAP[node.id])

这正是我正在寻找的（输出片段）：

Module
0
Import
18
alias
91
alias
91
alias
91
Import
18
alias

我现在的主要问题是提取 NODE_MAP[node.id]，因为 return 只能用于返回修改后的树。

Answer 1

ast.dump 显示node.__class__.__name__，所以我猜这是您要映射到数字的字符串，而数字又由NODE_LIST 中的索引确定。

class CustomNodeVisitor(ast.NodeVisitor):
    def visit(self, node):
        print(node.__class__.__name__)
        return ast.NodeVisitor.visit(self, node)