demand.py

import re
import sys

indent_raw_re = r'^(?P<indent>\s*)'
assign_raw_re = indent_raw_re + r'(?P<identifier>[\w\d_]+)\s*[+\-*/!^&]?='
trace_raw_re = assign_raw_re + r'.*?(?P<trace>\s+# = .*)'

indent_regex = re.compile(indent_raw_re)
identifier_regex = re.compile(assign_raw_re)
trace_regex = re.compile(trace_raw_re)


compose = lambda f: lambda g: lambda *a, **k: f(g(*a, **k))

def get_identifier(line):
    m = identifier_regex.match(line)
    if m:
        return m['identifier']

def strip_trace_line(line):
    m = trace_regex.match(line)
    if m:
        return line[:m.start('trace')]
    else:
        return line

def strip_traces(prog):
    return '\n'.join(strip_trace_line(line) for line in prog.split('\n'))

def is_quiet(line):
    mquots = "'''", '"""'
    return any(line.find(x) >= 0 for x in mquots) or line.strip().endswith(';')

@compose('\n'.join)
def tracify(prog):
    for n, line in enumerate(prog.split('\n')):
        identifier = get_identifier(line)
        if identifier and not is_quiet(line):
            yield line + " ; DEMAND({}, {})".format(n, identifier)
        else:
            yield line


def runtrace(prog):
    trprog = tracify(prog)
    traced_values = {}
    def tracer_func(n, x):
        if isinstance(x, float):
            r = '{:0.2f}'.format(x)
        else:
            r = repr(x)
        if n in traced_values:
            traced_values[n] += ', = ' + r
        else:
            traced_values[n] = ' # = ' + r
    noop = lambda *_, **__: None
    exec(trprog, {'DEMAND': tracer_func, 'print': noop})
    return traced_values


@compose('\n'.join)
def demand(rawprog):
    proglines = [strip_trace_line(line) for line in rawprog.split('\n')]
    prog = '\n'.join(proglines)
    values = runtrace(prog)
    for n, line in enumerate(proglines):
        yield line + values.get(n, '')


if __name__ == '__main__':
    filename = sys.argv[1] if len(sys.argv) > 1 else None
    strip = sys.argv[2] == '--strip' if len(sys.argv) > 2 else None
    if filename:
        with open(filename) as f:
            if strip:
                sys.stdout.write(strip_traces(f.read()))
            else:
                sys.stdout.write(demand(f.read()))

demand.vim

:nnoremap <f2> :let demandline=line('.')<cr>:w<cr>:%d<cr>:r !python3 demand.py %<cr>ggdd:exe demandline<cr>
nnoremap <f3> :let demandline=line('.')<cr>:w<cr>:%d<cr>:r !python3 demand.py % --strip<cr>ggdd:exe demandline<cr>

explanation

"""
x = 5
for i in range(5):
    y = i + 5
    continue
"""
>> tracify >>
"""
x = 5 ; DEMAND(1, x)
for i in range(5):
    y = i + 5 ; DEMAND(3, y)
    continue
"""
# Output:
"""
x = 5 # = 5
for i in range(5):
    y = i + 5 #  = 5, = 6, = 7, = 8, = 9
    continue
"""