python - 用注释注释 Python print() 输出

Question

给定一个包含 print() 语句的 Python 脚本，我希望能够运行该脚本并在每个语句后插入一条注释以显示每个语句的输出。为了演示，使用这个名为 example.py 的脚本:

a, b = 1, 2

print('a + b:', a + b)

c, d = 3, 4

print('c + d:', c + d)

期望的输出是:

a, b = 1, 2

print('a + b:', a + b)
# a + b: 3

c, d = 3, 4

print('c + d:', c + d)
# c + d: 7

这是我的尝试，适用于像上面这样的简单示例:

import sys
from io import StringIO

def intercept_stdout(func):
    "redirect stdout from a target function"
    def wrapper(*args, **kwargs):
        "wrapper function for intercepting stdout"
        # save original stdout
        original_stdout = sys.stdout

        # set up StringIO object to temporarily capture stdout
        capture_stdout = StringIO()
        sys.stdout = capture_stdout

        # execute wrapped function
        func(*args, **kwargs)

        # assign captured stdout to value
        func_output = capture_stdout.getvalue()

        # reset stdout
        sys.stdout = original_stdout

        # return captured value
        return func_output

    return wrapper


@intercept_stdout
def exec_target(name):
    "execute a target script"
    with open(name, 'r') as f:    
        exec(f.read())


def read_target(name):
    "read source code from a target script & return it as a list of lines"
    with open(name) as f:
        source = f.readlines()

    # to properly format last comment, ensure source ends in a newline
    if len(source[-1]) >= 1 and source[-1][-1] != '\n':
        source[-1] += '\n'

    return source


def annotate_source(target):
    "given a target script, return the source with comments under each print()"
    target_source = read_target(target)

    # find each line that starts with 'print(' & get indices in reverse order
    print_line_indices = [i for i, j in enumerate(target_source)
                              if len(j) > 6 and j[:6] == 'print(']
    print_line_indices.reverse()

    # execute the target script and get each line output in reverse order
    target_output = exec_target(target)
    printed_lines = target_output.split('\n')
    printed_lines.reverse()

    # iterate over the source and insert commented target output line-by-line
    annotated_source = []
    for i, line in enumerate(target_source):
        annotated_source.append(line)
        if print_line_indices and i == print_line_indices[-1]:
            annotated_source.append('# ' + printed_lines.pop() + '\n')
            print_line_indices.pop()

    # return new annotated source as a string
    return ''.join(annotated_source)


if __name__ == '__main__':
    target_script = 'example.py'
    with open('annotated_example.py', 'w') as f:
        f.write(annotate_source(target_script))

但是，对于包含跨多行的 print() 语句的脚本，以及跨多行的 print() 语句，它会失败' 在一行的开头。在最好的情况下，它甚至可以用于函数内的 print() 语句。举个例子:

print('''print to multiple lines, first line
second line
third line''')

print('print from partial line, first part') if True else 0

1 if False else print('print from partial line, second part')

print('print from compound statement, first part'); pass

pass; print('print from compound statement, second part')

def foo():
    print('bar')

foo()

理想情况下，输出应该是这样的:

print('''print to multiple lines, first line
second line
third line''')
# print to multiple lines, first line
# second line
# third line

print('print from partial line, first part') if True else 0
# print from partial line, first part

1 if False else print('print from partial line, second part')
# print from partial line, second part

print('print from compound statement, first part'); pass
# print from compound statement, first part

pass; print('print from compound statement, second part')
# print from compound statement, second part

def foo():
    print('bar')

foo()
# bar

但是上面的脚本像这样破坏了它:

print('''print to multiple lines, first line
# print to multiple lines, first line
second line
third line''')

print('print from partial line, first part') if True else 0
# second line

1 if False else print('print from partial line, second part')

print('print from compound statement, first part'); pass
# third line

pass; print('print from compound statement, second part')

def foo():
    print('bar')

foo()

什么方法可以使这个过程更加稳健？

Answer 1

您是否考虑过使用 inspect 模块？如果你愿意说你总是想要最上面的调用旁边的注释，并且你正在注释的文件足够简单，你可以获得合理的结果。以下是我的尝试，它覆盖了内置的 print 函数并查看堆栈跟踪以确定调用 print 的位置:

import inspect
import sys
from io import StringIO

file_changes = {}

def anno_print(old_print, *args, **kwargs):
    (frame, filename, line_number,
     function_name, lines, index) = inspect.getouterframes(inspect.currentframe())[-2]
    if filename not in file_changes:
        file_changes[filename] = {}
    if line_number not in file_changes[filename]:
        file_changes[filename][line_number] = []
    orig_stdout = sys.stdout
    capture_stdout = StringIO()
    sys.stdout = capture_stdout
    old_print(*args, **kwargs)
    output = capture_stdout.getvalue()
    file_changes[filename][line_number].append(output)
    sys.stdout = orig_stdout
    return

def make_annotated_file(old_source, new_source):
    changes = file_changes[old_source]
    old_source_F = open(old_source)
    new_source_F = open(new_source, 'w')
    content = old_source_F.readlines()
    for i in range(len(content)):
        line_num = i + 1
        new_source_F.write(content[i])
        if content[i][-1] != '\n':
            new_source_F.write('\n')
        if line_num in changes:
            for output in changes[line_num]:
                output = output[:-1].replace('\n', '\n#') + '\n'
                new_source_F.write("#" + output)
    new_source_F.close()



if __name__=='__main__':
    target_source = "foo.py"
    old_print = __builtins__.print
    __builtins__.print = lambda *args, **kwargs: anno_print(old_print, *args, **kwargs)
    with open(target_source) as f:
        code = compile(f.read(), target_source, 'exec')
        exec(code)
    __builtins__.print = old_print
    make_annotated_file(target_source, "foo_annotated.py")

如果我在以下文件“foo.py”上运行它:

def foo():
    print("a")
    print("b")

def cool():
    foo()
    print("c")

def doesnt_print():
    a = 2 + 3

print(1+2)
foo()
doesnt_print()
cool()

输出是“foo_annotated.py”:

def foo():
    print("a")
    print("b")

def cool():
    foo()
    print("c")

def doesnt_print():
    a = 2 + 3

print(1+2)
#3
foo()
#a
#b
doesnt_print()
cool()
#a
#b
#c