Add a script for interactively viewing traces

1 files changed, 148 insertions, 0 deletions

diff --git a/scripts/trace_viewer.py b/scripts/trace_viewer.py
new file mode 100644
index 0000000..8758187
--- /dev/null
+++ b/scripts/trace_viewer.py
@@ -0,0 +1,148 @@
+#!/usr/bin/python3
+
+import curses
+import sys
+
+from monkey.action import parse, expand
+from db.models import Problem, Solution
+
+if len(sys.argv) > 1:
+    pid = int(sys.argv[1])
+else:
+    # print all problem ids
+    print('problems:')
+    for problem in Problem.list():
+        print('  {}\t{}'.format(problem.id, problem.identifier))
+    print()
+    pid = int(input('enter problem id: '))
+attempts = sorted(Solution.filter(problem_id=pid), key=lambda s: s.codeq_user_id)
+
+def get_actions(trace):
+    try:
+        actions = parse(trace)
+        expand(actions)
+        return actions
+    except Exception as ex:
+        sys.stderr.write('Error parsing action log: {}\n'.format(ex))
+        return []
+
+def main(stdscr):
+    def draw_code(win, text, highlight_pos=None, highlight_color=0):
+        win.erase()
+        win.border()
+        y = x = 0
+        for n, char in enumerate(text + ' '):
+            attr = curses.color_pair(highlight_color if n == highlight_pos else 1)
+            if char == '\t':
+                char = ' '
+
+            if char != '\n':
+                win.addstr(y+1, x+1, char, attr)
+                x += 1
+            else:
+                # add an empty cell at end of each line for highlighting
+                win.addstr(y+1, x+1, ' ', attr)
+                x = 0
+                y += 1
+        win.refresh()
+
+    def draw_status(win, uid, action_idx):
+        win.erase()
+        text = 'usr {}, act {}/{}: {}'.format(uid, action_idx, len(actions),
+                                              actions[min(len(actions)-1, index)])
+        win.addstr(0, 0, text[:79])
+        win.refresh()
+
+    stdscr.clear()
+    if curses.has_colors():
+        curses.init_pair(1, curses.COLOR_WHITE, 0)
+        curses.init_pair(2, curses.COLOR_BLACK, curses.COLOR_GREEN)
+        curses.init_pair(3, curses.COLOR_WHITE, curses.COLOR_RED)
+    curses.init_color(0, 0, 0, 0)
+    curses.curs_set(0)
+
+    stdscr.attrset(curses.color_pair(1))
+    statuswin = stdscr.derwin(1, 80, 0, 0)
+    codewin = stdscr.derwin(22, 80, 1, 0)
+
+    attempt_idx = 0
+    if len(sys.argv) > 2:
+        uid = int(sys.argv[2])
+        for attempt_idx, a in enumerate(attempts):
+            if a.codeq_user_id == uid:
+                break
+    attempt = attempts[attempt_idx]
+    actions = get_actions(attempt.trace)
+    index = 0
+    code = ''
+
+    draw_status(statuswin, attempt.codeq_user_id, index)
+    draw_code(codewin, code, 0, 2)
+    while True:
+        c = stdscr.getch()
+        if 0 < c < 256:
+            c = chr(c)
+
+        if c == curses.KEY_HOME:
+            # go to first action
+            index = 0
+            code = ''
+        elif c == curses.KEY_LEFT:
+            # go to previous action
+            if index > 0:
+                index -= 1
+                action = actions[index]
+                code = action.unapply(code)
+        elif c == curses.KEY_RIGHT:
+            # go to next action
+            if index < len(actions):
+                action = actions[index]
+                code = action.apply(code)
+                index += 1
+        elif c == 'n':
+            # go to next solution
+            if attempt_idx < len(attempts)-1:
+                attempt_idx += 1
+                attempt = attempts[attempt_idx]
+                actions = get_actions(attempt.trace)
+                code = ''
+                index = 0
+        elif c == 'N':
+            # go to previous solution
+            if attempt_idx > 0:
+                attempt_idx -= 1
+                attempt = attempts[attempt_idx]
+                actions = get_actions(attempt.trace)
+                code = ''
+                index = 0
+        elif c == 'q':
+            # quit
+            break
+        elif c == 't':
+            # go to next 'test' action
+            if index < len(actions) and actions[index].type == 'test':
+                index += 1
+            while index < len(actions) and actions[index].type != 'test':
+                code = actions[index].apply(code)
+                index += 1
+        elif c == 'T':
+            # go to previous 'test' action
+            while index > 0 and actions[index-1].type != 'test':
+                code = actions[index-1].unapply(code)
+                index -= 1
+            if index > 0 and actions[index-1].type == 'test':
+                index -= 1
+
+        # if active action is 'remove' or 'insert', highlight the character
+        action = actions[index] if index < len(actions) else actions[-1]
+        if action.type in {'insert', 'remove'}:
+            highlight_pos = action.offset
+            highlight_color = 2 if action.type == 'insert' else 3
+        else:
+            highlight_pos = None
+            highlight_color = 0
+
+        draw_status(statuswin, attempt.codeq_user_id, index)
+        draw_code(codewin, code, highlight_pos, highlight_color)
+
+curses.wrapper(main)