Convert monkey.edits to use the new DB

1 files changed, 26 insertions, 22 deletions

diff --git a/monkey/edits.py b/monkey/edits.py
index 01e1ffd..0f386d5 100644
--- a/monkey/edits.py
+++ b/monkey/edits.py
@@ -237,31 +237,35 @@ def classify_edits(edits):
             changes[(before, after)] = cost
     return inserts, removes, changes
 
+
+# Extract edits and other data from existing traces for each problem.
 if __name__ == '__main__':
-    import os
     import pickle
-    import django
-
-    # Load django models.
-    os.environ['DJANGO_SETTINGS_MODULE'] = 'webmonkey.settings'
-    django.setup()
-    from django.contrib.auth.models import User
-    from tutor.models import Attempt, Problem
-
-    edits = {}
-    submissions = {}
-    queries = {}
-    names = {}
-    for problem in Problem.objects.all():
-        print(problem.name)
-        pid = problem.pk
-        attempts = Attempt.objects.filter(problem=problem, done=True) \
-                                  .exclude(user__groups=None)
-        traces = [a.trace for a in attempts]
+    from db.models import Problem, Solution
+
+    # Ignore traces from these users.
+    ignored_users = [
+        1, # admin
+        231, # test
+        360, # test2
+        358, # sasha
+    ]
+
+    edits, submissions, queries, names = {}, {}, {}, {}
+    for problem in Problem.list():
+        pid = problem.id
+        solutions = Solution.filter(problem_id=pid, done=True)
+        traces = [s.trace for s in solutions if s.codeq_user_id not in ignored_users]
+        if not traces:
+            print('No traces for {}'.format(problem.name))
+            continue
+
+        print('Analyzing traces for {}… '.format(problem.name), end='', flush=True)
         try:
             edits[pid], submissions[pid], queries[pid], names[pid] = get_edits_from_traces(traces)
-        except:
-            pass
+            print('{} edits, {} submissions, {} queries, {} names'.format(
+                len(edits[pid]), len(submissions[pid]), len(queries[pid]), len(names[pid])))
+        except Exception as ex:
+            print('error: {}'.format(ex))
 
     pickle.dump((edits, submissions, queries, names), open('edits.pickle', 'wb'))
-