From 07ca340c80839339ba82b8dd053d18aa522ed169 Mon Sep 17 00:00:00 2001
From: Ethan Dalool <ethan@voussoir.net>
Date: Sun, 28 Oct 2018 22:56:05 -0700
Subject: [PATCH] Add TSDB.for_submission to remove some duplicate effort.

---
 timesearch/commentaugment.py  | 21 +++++++++++----------
 timesearch/offline_reading.py |  7 ++++---
 timesearch/tsdb.py            |  8 ++++++++
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/timesearch/commentaugment.py b/timesearch/commentaugment.py
index 9c9b09e..a9d12e5 100644
--- a/timesearch/commentaugment.py
+++ b/timesearch/commentaugment.py
@@ -21,18 +21,19 @@ def commentaugment(
 
     common.login()
 
-    if specific_submission is not None:
-        subreddit = common.subreddit_for_submission(specific_submission).display_name
+    if specific_submission:
+        (database, subreddit) = tsdb.TSDB.for_submission(specific_submission, do_create=True, fix_name=True)
+        specific_submission = common.t3_prefix(specific_submission)[3:]
+        specific_submission = common.r.submission(specific_submission)
+        database.insert(specific_submission)
+
+    elif subreddit:
+        (database, subreddit) = tsdb.TSDB.for_subreddit(subreddit, do_create=False, fix_name=True)
 
-    if subreddit:
-        do_create = specific_submission is not None
-        (database, subreddit) = tsdb.TSDB.for_subreddit(subreddit, do_create=do_create, fix_name=True)
     else:
         (database, username) = tsdb.TSDB.for_user(username, do_create=False, fix_name=True)
-    cur = database.sql.cursor()
 
-    if specific_submission is not None:
-        database.insert(specific_submission_obj)
+    cur = database.sql.cursor()
 
     if lower is None:
         lower = 0
@@ -67,7 +68,7 @@ def commentaugment(
         lower = 0
 
     if specific_submission:
-        comments = pushshift.get_comments_from_submission(specific_submission_obj)
+        comments = pushshift.get_comments_from_submission(specific_submission)
     elif subreddit:
         comments = pushshift.get_comments_from_subreddit(subreddit, lower=lower, upper=upper)
     elif username:
@@ -92,7 +93,7 @@ def commentaugment(
             set augmented_at = ?
             WHERE idstr == ?
         '''
-        bindings = [common.get_now(), specific_submission_obj.fullname]
+        bindings = [common.get_now(), specific_submission.fullname]
         cur.execute(query, bindings)
         database.sql.commit()
 
diff --git a/timesearch/offline_reading.py b/timesearch/offline_reading.py
index a794d78..7c500b6 100644
--- a/timesearch/offline_reading.py
+++ b/timesearch/offline_reading.py
@@ -227,11 +227,12 @@ def html_from_database(subreddit=None, username=None, specific_submission=None):
     if username and specific_submission:
         raise exceptions.NotExclusive(['username', 'specific_submission'])
 
-    if specific_submission is not None:
-        subreddit = common.subreddit_for_submission(specific_submission).display_name
+    if specific_submission:
+        database = tsdb.TSDB.for_submission(specific_submission, do_create=False)
 
-    if subreddit:
+    elif subreddit:
         database = tsdb.TSDB.for_subreddit(subreddit, do_create=False)
+
     else:
         database = tsdb.TSDB.for_user(username, do_create=False)
 
diff --git a/timesearch/tsdb.py b/timesearch/tsdb.py
index e0d68f7..a45902b 100644
--- a/timesearch/tsdb.py
+++ b/timesearch/tsdb.py
@@ -215,6 +215,14 @@ class TSDB:
             return (database, name_from_path(name))
         return database
 
+    @classmethod
+    def for_submission(cls, submission_id, fix_name=False, *args, **kwargs):
+        subreddit = common.subreddit_for_submission(submission_id)
+        database = cls.for_subreddit(subreddit, *args, **kwargs)
+        if fix_name:
+            return (database, subreddit.display_name)
+        return database
+
     @classmethod
     def for_subreddit(cls, name, do_create=True, fix_name=False):
         if isinstance(name, common.praw.models.Subreddit):