沿执行路径收集python源代码注释

Question

例如，我有以下 python 函数：

def func(x):
    """Function docstring."""

    result = x + 1
    if result > 0:
        # comment 2
        return result
    else:
        # comment 3
        return -1 * result

而且我想要一些函数来打印沿执行路径遇到的所有函数文档字符串和注释，例如

> trace(func(2))
Function docstring.
Comment 2
3

事实上，我试图实现的是提供一些关于如何计算结果的评论。

可以用什么？据我所知，AST 不会在树中保留评论。

Answer 1

我认为这是一个有趣的挑战，所以我决定试一试。这是我想出的：

import ast
import inspect
import re
import sys
import __future__

if sys.version_info >= (3,5):
    ast_Call = ast.Call
else:
    def ast_Call(func, args, keywords):
        """Compatibility wrapper for ast.Call on Python 3.4 and below.
        Used to have two additional fields (starargs, kwargs)."""
        return ast.Call(func, args, keywords, None, None)

COMMENT_RE = re.compile(r'^(\s*)#\s?(.*)$')

def convert_comment_to_print(line):
    """If `line` contains a comment, it is changed into a print
    statement, otherwise nothing happens. Only acts on full-line comments,
    not on trailing comments. Returns the (possibly modified) line."""
    match = COMMENT_RE.match(line)
    if match:
        return '{}print({!r})\n'.format(*match.groups())
    else:
        return line

def convert_docstrings_to_prints(syntax_tree):
    """Walks an AST and changes every docstring (i.e. every expression
    statement consisting only of a string) to a print statement.
    The AST is modified in-place."""
    ast_print = ast.Name('print', ast.Load())
    nodes = list(ast.walk(syntax_tree))
    for node in nodes:
        for bodylike_field in ('body', 'orelse', 'finalbody'):
            if hasattr(node, bodylike_field):
                for statement in getattr(node, bodylike_field):
                    if (isinstance(statement, ast.Expr) and
                            isinstance(statement.value, ast.Str)):
                        arg = statement.value
                        statement.value = ast_Call(ast_print, [arg], [])

def get_future_flags(module_or_func):
    """Get the compile flags corresponding to the features imported from
    __future__ by the specified module, or by the module containing the
    specific function. Returns a single integer containing the bitwise OR
    of all the flags that were found."""
    result = 0
    for feature_name in __future__.all_feature_names:
        feature = getattr(__future__, feature_name)
        if (hasattr(module_or_func, feature_name) and
                getattr(module_or_func, feature_name) is feature and
                hasattr(feature, 'compiler_flag')):
            result |= feature.compiler_flag
    return result

def eval_function(syntax_tree, func_globals, filename, lineno, compile_flags,
        *args, **kwargs):
    """Helper function for `trace`. Execute the function defined by
    the given syntax tree, and return its return value."""
    func = syntax_tree.body[0]
    func.decorator_list.insert(0, ast.Name('_trace_exec_decorator', ast.Load()))
    ast.increment_lineno(syntax_tree, lineno-1)
    ast.fix_missing_locations(syntax_tree)
    code = compile(syntax_tree, filename, 'exec', compile_flags, True)
    result = [None]
    def _trace_exec_decorator(compiled_func):
        result[0] = compiled_func(*args, **kwargs)
    func_locals = {'_trace_exec_decorator': _trace_exec_decorator}
    exec(code, func_globals, func_locals)
    return result[0]

def trace(func, *args, **kwargs):
    """Run the given function with the given arguments and keyword arguments,
    and whenever a docstring or (whole-line) comment is encountered,
    print it to stdout."""
    filename = inspect.getsourcefile(func)
    lines, lineno = inspect.getsourcelines(func)
    lines = map(convert_comment_to_print, lines)
    modified_source = ''.join(lines)
    compile_flags = get_future_flags(func)
    syntax_tree = compile(modified_source, filename, 'exec',
            ast.PyCF_ONLY_AST | compile_flags, True)
    convert_docstrings_to_prints(syntax_tree)
    return eval_function(syntax_tree, func.__globals__,
            filename, lineno, compile_flags, *args, **kwargs)

它有点长，因为我试图涵盖最重要的情况，并且代码可能不是最易读的，但我希望它足够好理解。

这个怎么运作：

1. 首先，使用 阅读函数的源代码inspect.getsourcelines。（警告：inspect不适用于以交互方式定义的函数。如果您需要，也许可以dill改用，请参阅此答案。）
2. 搜索看起来像注释的行，并用打印语句替换它们。（现在只替换整行注释，但如果需要，将其扩展到尾随注释应该不难。）
3. 将源代码解析为 AST。
4. 走 AST 并用打印语句替换所有文档字符串。
5. 编译 AST。
6. 执行 AST。这一步和上一步包含一些技巧来尝试重建函数最初定义的上下文（例如全局变量、__future__导入、异常回溯的行号）。此外，由于只执行源代码只会重新定义函数而不是调用它，因此我们使用一个简单的装饰器来修复它。

它适用于 Python 2 和 3（至少在下面的测试中，我在 2.7 和 3.6 中运行）。

要使用它，只需执行以下操作：

result = trace(func, 2)   # result = func(2)

这是我在编写代码时使用的一个稍微复杂的测试：

#!/usr/bin/env python

from trace_comments import trace
from dateutil.easter import easter, EASTER_ORTHODOX

def func(x):
    """Function docstring."""

    result = x + 1
    if result > 0:
        # comment 2
        return result
    else:
        # comment 3
        return -1 * result

if __name__ == '__main__':
    result1 = trace(func, 2)
    print("result1 = {}".format(result1))

    result2 = trace(func, -10)
    print("result2 = {}".format(result2))

    # Test that trace() does not permanently replace the function
    result3 = func(42)
    print("result3 = {}".format(result3))

    print("-----")
    print(trace(easter, 2018))

    print("-----")
    print(trace(easter, 2018, EASTER_ORTHODOX))