Python write both commands and their output to a file

Here's a script I just wrote which quite comprehensively captures printed output and prints it alongside the code, no matter how it's printed or how much is printed in one go. It uses the ast module to parse the Python source, executes the program one statement at a time (kind of as if it was fed to the REPL), then prints the output from each statement. Python 3.6+ (but easily modified for e.g. Python 2.x):

import ast
import sys

if len(sys.argv) < 2:
    print(f"Usage: {sys.argv[0]} <script.py> [args...]")
    exit(1)

# Replace stdout so we can mix program output and source code cleanly
real_stdout = sys.stdout
class FakeStdout:
    ''' A replacement for stdout that prefixes # to every line of output, so it can be mixed with code. '''
    def __init__(self, file):
        self.file = file
        self.curline = ''

    def _writerow(self, row):
        self.file.write('# ')
        self.file.write(row)
        self.file.write('\n')

    def write(self, text):
        if not text:
            return
        rows = text.split('\n')
        self.curline += rows.pop(0)
        if not rows:
            return
        for row in rows:
            self._writerow(self.curline)
            self.curline = row

    def flush(self):
        if self.curline:
            self._writerow(self.curline)
            self.curline = ''

sys.stdout = FakeStdout(real_stdout)

class EndLineFinder(ast.NodeVisitor):
    ''' This class functions as a replacement for the somewhat unreliable end_lineno attribute.

    It simply finds the largest line number among all child nodes. '''

    def __init__(self):
        self.max_lineno = 0

    def generic_visit(self, node):
        if hasattr(node, 'lineno'):
            self.max_lineno = max(self.max_lineno, node.lineno)
        ast.NodeVisitor.generic_visit(self, node)

# Pretend the script was called directly
del sys.argv[0]

# We'll walk each statement of the file and execute it separately.
# This way, we can place the output for each statement right after the statement itself.
filename = sys.argv[0]
source = open(filename, 'r').read()
lines = source.split('\n')
module = ast.parse(source, filename)
env = {'__name__': '__main__'}

prevline = 0
endfinder = EndLineFinder()

for stmt in module.body:
    # note: end_lineno will be 1-indexed (but it's always used as an endpoint, so no off-by-one errors here)
    endfinder.visit(stmt)
    end_lineno = endfinder.max_lineno
    for line in range(prevline, end_lineno):
        print(lines[line], file=real_stdout)
    prevline = end_lineno
    # run a one-line "module" containing only this statement
    exec(compile(ast.Module([stmt]), filename, 'exec'), env)
    # flush any incomplete output (FakeStdout is "line-buffered")
    sys.stdout.flush()

Here's a test script:

print(3); print(4)
print(5)

if 1:
    print(6)

x = 3
for i in range(6):
    print(x + i)

import sys
sys.stdout.write('I love Python')

import pprint
pprint.pprint({'a': 'b', 'c': 'd'}, width=5)

and the result:

print(3); print(4)
# 3
# 4
print(5)
# 5

if 1:
    print(6)
# 6

x = 3
for i in range(6):
    print(x + i)
# 3
# 4
# 5
# 6
# 7
# 8

import sys
sys.stdout.write('I love Python')
# I love Python

import pprint
pprint.pprint({'a': 'b', 'c': 'd'}, width=5)
# {'a': 'b',
#  'c': 'd'}