Collect python source code comments along execution path

Question

E.g. I've got the following python function:

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

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

And I want to have some function that would print all function docstrings and comments that are met along the execution path, e.g.

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

In fact what I try to achieve is to provide some comments how the result has been calculated.

What could be used? AST as far as I understand does not keep comment in the tree.

Answer 1

I thought this was an interesting challenge, so I decided to give it a try. Here is what I came up with:

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)

It is a bit long because I tried to cover most important cases, and the code might not be the most readable, but I hope it is nice enough to follow.

How it works:

1. First, read the function's source code using inspect.getsourcelines. (Warning: inspect does not work for functions that were defined interactively. If you need that, maybe you can use dill instead, see this answer.)
2. Search for lines that look like comments, and replace them with print statements. (Right now only whole-line comments are replaced, but it shouldn't be difficult to extend that to trailing comments if desired.)
3. Parse the source code into an AST.
4. Walk the AST and replace all docstrings with print statements.
5. Compile the AST.
6. Execute the AST. This and the previous step contain some trickery to try to reconstruct the context that the function was originally defined in (e.g. globals, __future__ imports, line numbers for exception tracebacks). Also, since just executing the source would only re-define the function and not call it, we fix that with a simple decorator.

It works in Python 2 and 3 (at least with the tests below, which I ran in 2.7 and 3.6).

To use it, simply do:

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

Here is a slightly more elaborate test that I used while writing the code:

#!/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))