保留注释的Python AST

Question

我可以通过下面的代码获取没有注释的抽象语法树（AST）：

import ast
module = ast.parse(open('/path/to/module.py').read())

你能举个例子，说明如何获取保留注释（和空白）的抽象语法树吗？

Answer 1

在写任何类型的Python代码美化工具、pep-8检查器等时，这个问题自然会出现。在这种情况下，你实际上是在进行源代码到源代码的转换，你希望输入是人写的，不仅希望输出是人能读懂的，还希望它能：

保留所有注释，准确放在原来的位置。
输出字符串的准确拼写，包括文档字符串，和原始内容一致。

用ast模块来做到这一点并不简单。可以说这是API的一个缺陷，但似乎没有简单的方法可以扩展API来轻松实现第1和第2点。

Andrei建议同时使用ast和tokenize，这个方法非常聪明。当我在写一个Python转Coffeescript的转换器时，我也有过类似的想法，但代码并不简单。

在py2cs.py文件中，从第1305行开始的TokenSync（ts）类负责协调基于token的数据和ast遍历之间的通信。给定源字符串s，TokenSync类会对s进行分词，并初始化内部数据结构，支持几个接口方法：

ts.leading_lines(node)：返回前面的注释和空行的列表。

ts.trailing_comment(node)：返回节点的尾部注释字符串（如果有的话）。

ts.sync_string(node)：返回给定节点的字符串拼写。

对于ast访问者来说，使用这些方法是直接的，但有点笨拙。以下是来自py2cs.py中CoffeeScriptTraverser（cst）类的一些示例：

def do_Str(self, node):
    '''A string constant, including docstrings.'''
    if hasattr(node, 'lineno'):
        return self.sync_string(node)

只要按照源代码中出现的顺序访问ast.Str节点，这个方法就能正常工作。这在大多数遍历中是自然而然发生的。

这里是ast.If访问者。它展示了如何使用ts.leading_lines和ts.trailing_comment：

def do_If(self, node):

    result = self.leading_lines(node)
    tail = self.trailing_comment(node)
    s = 'if %s:%s' % (self.visit(node.test), tail)
    result.append(self.indent(s))
    for z in node.body:
        self.level += 1
        result.append(self.visit(z))
        self.level -= 1
    if node.orelse:
        tail = self.tail_after_body(node.body, node.orelse, result)
        result.append(self.indent('else:' + tail))
        for z in node.orelse:
            self.level += 1
            result.append(self.visit(z))
            self.level -= 1
    return ''.join(result)

ts.tail_after_body方法弥补了没有表示'else'子句的ast节点这一事实。这并不复杂，但也不太好看：

def tail_after_body(self, body, aList, result):
    '''
    Return the tail of the 'else' or 'finally' statement following the given body.
    aList is the node.orelse or node.finalbody list.
    '''
    node = self.last_node(body)
    if node:
        max_n = node.lineno
        leading = self.leading_lines(aList[0])
        if leading:
            result.extend(leading)
            max_n += len(leading)
        tail = self.trailing_comment_at_lineno(max_n + 1)
    else:
        tail = '\n'
    return tail

注意，cst.tail_after_body只是调用了ts.tail_after_body。

总结

TokenSync类封装了将基于token的数据提供给ast遍历代码所涉及的大部分复杂性。使用TokenSync类是直接的，但所有Python语句的ast访问者（以及ast.Str）必须包含对ts.leading_lines、ts.trailing_comment和ts.sync_string的调用。此外，ts.tail_after_body这个小技巧是处理“缺失”ast节点所必需的。

简而言之，这段代码工作得很好，但有点笨拙。

@Andrei：你的简短回答可能暗示你知道更优雅的方法。如果是这样，我很想看看。

Edward K. Ream

Answer 2

一个包含格式、注释等信息的抽象语法树叫做完整语法树。

redbaron 可以做到这一点。你可以通过 pip install redbaron 来安装它，然后试试以下代码。

import redbaron

with open("/path/to/module.py", "r") as source_code:
    red = redbaron.RedBaron(source_code.read())

print (red.fst())

Answer 3

ast模块不包含注释。虽然tokenize模块可以提取注释，但它不提供其他程序结构的信息。

保留注释的Python AST

8 个回答

撰写回答