adding separate database class

src/word_stats.py

@@ -1,34 +1,30 @@
 from collections import defaultdict, Counter
 
 from progress_bar import progress
-import sqlite3
 
 
 class WordStats:
-    def __init__(self, lemma_features):
+    def __init__(self, lemma_features, db):
         self.lemma_features = lemma_features
+        self.db = db
 
         self.all_words = 0
         self.memoized_msd_merges = {}
 
-        with open("sqlite.db", 'w') as fp:
-            fp.write("")
-
-        self.db = sqlite3.connect('sqlite.db')
-        self.db.execute("""CREATE TABLE UniqWords (
+        self.db.init("""CREATE TABLE UniqWords (
             uw_id INTEGER PRIMARY KEY, 
             lemma varchar(64), 
             msd varchar(16), 
             text varchar(64), 
             frequency int
             )""")
-        self.db.execute("CREATE TABLE CommonMsd (lemma varchar(64), msd varchar(16))")
-        self.db.execute("CREATE TABLE WordCount (lemma varchar(64), msd0 char, frequency int)")
+        self.db.init("CREATE TABLE CommonMsd (lemma varchar(64), msd varchar(16))")
+        self.db.init("CREATE TABLE WordCount (lemma varchar(64), msd0 char, frequency int)")
 
-        self.db.execute("CREATE INDEX lemma_msd_text_on_uw ON UniqWords (lemma, msd, text)")
-        self.db.execute("CREATE INDEX lemma_on_uw ON UniqWords (lemma)")
-        self.db.execute("CREATE INDEX lemma_on_cm ON CommonMsd (lemma)")
-        self.db.execute("CREATE INDEX lemma_msd0_on_wc ON WordCount (lemma, msd0)")
+        self.db.init("CREATE INDEX lemma_msd_text_on_uw ON UniqWords (lemma, msd, text)")
+        self.db.init("CREATE INDEX lemma_on_uw ON UniqWords (lemma)")
+        self.db.init("CREATE INDEX lemma_on_cm ON CommonMsd (lemma)")
+        self.db.init("CREATE INDEX lemma_msd0_on_wc ON WordCount (lemma, msd0)")
 
     def add_words(self, words):
         for w in progress(words, "adding-words"):