python - 沿执行路径收集 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 代替，请参阅 this answer。)
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))