Add support for pickles to trace viewer

1 files changed, 8 insertions, 2 deletions

diff --git a/scripts/trace_viewer.py b/scripts/trace_viewer.py
index ebf75b7..6d5e1f4 100644
--- a/scripts/trace_viewer.py
+++ b/scripts/trace_viewer.py
@@ -15,13 +15,19 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import curses
+import pickle
 import sys
 
 from monkey.action import parse, expand
 from db.models import Problem, Solution
 
 if len(sys.argv) > 1:
-    pid = int(sys.argv[1])
+    if sys.argv[1].endswith('.pickle'):
+        with open(sys.argv[1], 'rb') as f:
+            attempts = pickle.load(f)
+    else:
+        pid = int(sys.argv[1])
+        attempts = sorted(Solution.filter(problem_id=pid), key=lambda s: s.codeq_user_id)
 else:
     # print all problem ids
     print('problems:')
@@ -29,7 +35,7 @@ else:
         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)
+    attempts = sorted(Solution.filter(problem_id=pid), key=lambda s: s.codeq_user_id)
 
 def get_actions(trace):
     try: