Question

E.g。我有以下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)

这有点长，因为我试图涵盖最重要的案例，而且代码可能不是最易读的，但我希望它足够好了。

工作原理：

首先，使用inspect.getsourcelines阅读函数的源代码。（警告：inspect不适用于以交互方式定义的功能。如果您需要，可以使用dill，请参阅this answer。）
搜索看起来像评论的行，并将其替换为print语句。（现在只替换整行注释，但如果需要，不应该将其扩展到尾随注释。）
将源代码解析为AST。
使用print语句转换AST并替换所有文档字符串。
编译AST。
执行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))

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

1 个答案: