Move the DBEntry class to tsdb.

timesearch/offline_reading.py

@@ -98,24 +98,6 @@ HTML_SUBMISSION = '''
 '''.strip()
 
 
-class DBEntry:
-    def __init__(self, fetch):
-        if fetch[1].startswith('t3_'):
-            columns = tsdb.SQL_SUBMISSION_COLUMNS
-            self.object_type = 'submission'
-        else:
-            columns = tsdb.SQL_COMMENT_COLUMNS
-            self.object_type = 'comment'
-
-        self.id = None
-        self.idstr = None
-        for (index, attribute) in enumerate(columns):
-            setattr(self, attribute, fetch[index])
-
-    def __repr__(self):
-        return 'DBEntry(\'%s\')' % self.id
-
-
 class TreeNode:
     def __init__(self, identifier, data, parent=None):
         assert isinstance(identifier, str)
@@ -352,8 +334,8 @@ def tree_from_submission(submission_dbrow, comments_dbrows):
     Given the sqlite data for a submission and all of its comments,
     return a tree with the submission id as the root
     '''
-    submission = DBEntry(submission_dbrow)
+    submission = tsdb.DBEntry(submission_dbrow)
-    comments = [DBEntry(c) for c in comments_dbrows]
+    comments = [tsdb.DBEntry(c) for c in comments_dbrows]
     comments.sort(key=lambda x: x.created)
 
     print('Building tree for %s (%d comments)' % (submission.idstr, len(comments)))

timesearch/tsdb.py

@@ -141,6 +141,29 @@ SQL_COMMENT = {key:index for (index, key) in enumerate(SQL_COMMENT_COLUMNS)}
 SUBMISSION_TYPES = (common.praw.models.Submission, pushshift.DummySubmission)
 COMMENT_TYPES = (common.praw.models.Comment, pushshift.DummyComment)
 
+
+class DBEntry:
+    '''
+    This class converts a tuple row from the database into an object so that
+    you can access the attributes with dot notation.
+    '''
+    def __init__(self, dbrow):
+        if dbrow[1].startswith('t3_'):
+            columns = SQL_SUBMISSION_COLUMNS
+            self.object_type = 'submission'
+        else:
+            columns = SQL_COMMENT_COLUMNS
+            self.object_type = 'comment'
+
+        self.id = None
+        self.idstr = None
+        for (index, attribute) in enumerate(columns):
+            setattr(self, attribute, dbrow[index])
+
+    def __repr__(self):
+        return 'DBEntry(\'%s\')' % self.id
+
+
 class TSDB:
     def __init__(self, filepath, do_create=True):
         self.filepath = pathclass.Path(filepath)