storing depth used for parse in db, no need to specify it when generating

Author: codebox

db.py

@@ -1,38 +1,55 @@
 
 class Db:
-	def __init__(self, depth, conn, sql):
-		self.depth  = depth
+	DEPTH_PARAM_NAME = 'depth'
+	
+	def __init__(self, conn, sql):
 		self.conn   = conn 
 		self.cursor = conn.cursor()
 		self.sql    = sql
+		self.depth  = None
 
-		self.cursor.execute(self.sql.create_table_sql())
-		self.cursor.execute(self.sql.create_index_sql())
+	def setup(self, depth):
+		self.depth = depth
+		self.cursor.execute(self.sql.create_word_table_sql(depth))
+		self.cursor.execute(self.sql.create_index_sql(depth))
+		self.cursor.execute(self.sql.create_param_table_sql())
+		self.cursor.execute(self.sql.set_param_sql(), (self.DEPTH_PARAM_NAME, depth))
 
 	def _get_word_list_count(self, word_list):
-		if len(word_list) != self.depth:
-			raise ValueError('Expected %s words in list but found %s' % (self.depth, len(word_list)))
+		if len(word_list) != self.get_depth():
+			raise ValueError('Expected %s words in list but found %s' % (self.get_depth(), len(word_list)))
 
-		self.cursor.execute(self.sql.select_count_for_words_sql(), word_list)
+		self.cursor.execute(self.sql.select_count_for_words_sql(self.get_depth()), word_list)
 		r = self.cursor.fetchone()
 		if r:
 			return r[0]
 		else:
 			return 0
 
+	def get_depth(self):
+		if self.depth == None:
+			self.cursor.execute(self.sql.get_param_sql(), (self.DEPTH_PARAM_NAME,))
+			r = self.cursor.fetchone()
+			if r:
+				self.depth = int(r[0])
+			else:
+				raise ValueError('No depth value found in database, db does not seem to have been created by this utility')
+			
+		return self.depth
+		
 	def add_word(self, word_list):
 		count = self._get_word_list_count(word_list)
 		if count:
-			self.cursor.execute(self.sql.update_count_for_words_sql(), [count + 1] + word_list)
+			self.cursor.execute(self.sql.update_count_for_words_sql(self.get_depth()), [count + 1] + word_list)
 		else:
-			self.cursor.execute(self.sql.insert_row_for_words_sql(), word_list + [1])
+			self.cursor.execute(self.sql.insert_row_for_words_sql(self.get_depth()), word_list + [1])
 
 	def commit(self):
 		self.conn.commit()
 
 	def get_word_count(self, word_list):
 		counts = {}
-		sql = self.sql.select_words_and_counts_sql()
+		sql = self.sql.select_words_and_counts_sql(self.get_depth())
 		for row in self.cursor.execute(sql, word_list):
 			counts[row[0]] = row[1]
 

gen.py

@@ -17,7 +17,8 @@ def _get_next_word(self, word_list):
 				return w
 		assert False
 
-	def generate(self, depth, word_separator):
+	def generate(self, word_separator):
+		depth = self.db.get_depth()
 		sentence = [Parser.SENTENCE_START_SYMBOL] * (depth - 1)
 		end_symbol = [Parser.SENTENCE_END_SYMBOL] * (depth - 1)
 

markov.py

@@ -7,32 +7,38 @@
 import sqlite3
 import codecs
 
-SENTENCE_SEPARATOR = '\n'
-WORD_SEPARATOR = ''
+SENTENCE_SEPARATOR = '.'
+WORD_SEPARATOR = ' '
 
 if __name__ == '__main__':
 	args = sys.argv
-	usage = 'Usage: %s (parse <name> <depth> <path to txt file>|gen <name> <depth> <count>)' % (args[0], )
+	usage = 'Usage: %s (parse <name> <depth> <path to txt file>|gen <name> <count>)' % (args[0], )
 
-	if (len(args) != 5):
+	if (len(args) < 3):
 		raise ValueError(usage)
 
 	mode  = args[1]
 	name  = args[2]
-	depth = int(args[3])
-	sql   = Sql(depth)
-	db    = Db(depth, sqlite3.connect(name + '.db'), sql)
 
 	if mode == 'parse':
+		if (len(args) != 5):
+			raise ValueError(usage)
+		
+		depth = int(args[3])
 		file_name = args[4]
+		
+		db = Db(sqlite3.connect(name + '.db'), Sql())
+		db.setup(depth)
+		
 		txt = codecs.open(file_name, 'r', 'utf-8').read()
-		Parser(name, db, SENTENCE_SEPARATOR, WORD_SEPARATOR).parse(depth, txt)
+		Parser(name, db, SENTENCE_SEPARATOR, WORD_SEPARATOR).parse(txt)
 
 	elif mode == 'gen':
-		count = int(args[4])
+		count = int(args[3])
+		db = Db(sqlite3.connect(name + '.db'), Sql())
 		generator = Generator(name, db, Rnd())
 		for i in range(0, count):
-			print generator.generate(depth, WORD_SEPARATOR)
+			print generator.generate(WORD_SEPARATOR)
 
 	else:
 		raise ValueError(usage)

parse.py

@@ -1,4 +1,5 @@
 import sys
+import re
 
 class Parser:
 	SENTENCE_START_SYMBOL = '^'
@@ -9,12 +10,16 @@ def __init__(self, name, db, sentence_split_char = '\n', word_split_char = ''):
 		self.db   = db
 		self.sentence_split_char = sentence_split_char
 		self.word_split_char = word_split_char
+		self.whitespace_regex = re.compile('\s+')
 
-	def parse(self, depth, txt):
+	def parse(self, txt):
+		depth = self.db.get_depth()
 		sentences = txt.split(self.sentence_split_char)
 		i = 0
 
 		for sentence in sentences:
+			sentence = self.whitespace_regex.sub(" ", sentence).strip()
+
 			list_of_words = None
 			if self.word_split_char:
 				list_of_words = sentence.split(self.word_split_char)

sql.py

@@ -1,43 +1,54 @@
 class Sql:
     WORD_COL_NAME_PREFIX = 'word'
     COUNT_COL_NAME       = 'count'
-    TABLE_NAME           = 'word'
+    WORD_TABLE_NAME      = 'word'
     INDEX_NAME           = 'i_word'
+    PARAM_TABLE_NAME     = 'param'
+    KEY_COL_NAME         = 'name'
+    VAL_COL_NAME         = 'value'
 
-    def __init__(self, column_count):
-        self.column_count = int(column_count)
-        if self.column_count < 2:
+    def _check_column_count(self, count):
+        if count < 2:
             raise ValueError('Invalid column_count value, must be >= 2')
 
-    def _make_column_name_list(self):
-        return ', '.join([self.WORD_COL_NAME_PREFIX + str(n) for n in range(1, self.column_count + 1)])
+    def _make_column_name_list(self, column_count):
+        return ', '.join([self.WORD_COL_NAME_PREFIX + str(n) for n in range(1, column_count + 1)])
 
-    def _make_column_names_and_placeholders(self, col_count):
-        return ' AND '.join(['%s%s=?' % (self.WORD_COL_NAME_PREFIX, n) for n in range(1, col_count + 1)])
+    def _make_column_names_and_placeholders(self, column_count):
+        return ' AND '.join(['%s%s=?' % (self.WORD_COL_NAME_PREFIX, n) for n in range(1, column_count + 1)])
 
-    def create_table_sql(self):
-        return 'CREATE TABLE IF NOT EXISTS %s (%s, %s)' % (self.TABLE_NAME, self._make_column_name_list(), self.COUNT_COL_NAME)
+    def create_word_table_sql(self, column_count):
+        return 'CREATE TABLE IF NOT EXISTS %s (%s, %s)' % (self.WORD_TABLE_NAME, self._make_column_name_list(column_count), self.COUNT_COL_NAME)
 
-    def create_index_sql(self):
-        return 'CREATE INDEX IF NOT EXISTS %s ON %s (%s)' % (self.INDEX_NAME, self.TABLE_NAME, self._make_column_name_list())
+    def create_param_table_sql(self):
+        return 'CREATE TABLE IF NOT EXISTS %s (%s, %s)' % (self.PARAM_TABLE_NAME, self.KEY_COL_NAME, self.VAL_COL_NAME)
 
-    def select_count_for_words_sql(self):
-        return 'SELECT %s FROM %s WHERE %s' % (self.COUNT_COL_NAME, self.TABLE_NAME, self._make_column_names_and_placeholders(self.column_count)) 
+    def set_param_sql(self):
+        return 'INSERT INTO %s (%s, %s) VALUES (?, ?)' % (self.PARAM_TABLE_NAME, self.KEY_COL_NAME, self.VAL_COL_NAME)
 
-    def update_count_for_words_sql(self):
-        return 'UPDATE %s SET %s=? WHERE %s' % (self.TABLE_NAME, self.COUNT_COL_NAME, self._make_column_names_and_placeholders(self.column_count)) 
+    def get_param_sql(self):
+        return 'SELECT %s FROM %s WHERE %s=?' % (self.VAL_COL_NAME, self.PARAM_TABLE_NAME, self.KEY_COL_NAME)
+
+    def create_index_sql(self, column_count):
+        return 'CREATE INDEX IF NOT EXISTS %s ON %s (%s)' % (self.INDEX_NAME, self.WORD_TABLE_NAME, self._make_column_name_list(column_count))
+    
+    def select_count_for_words_sql(self, column_count):
+        return 'SELECT %s FROM %s WHERE %s' % (self.COUNT_COL_NAME, self.WORD_TABLE_NAME, self._make_column_names_and_placeholders(column_count)) 
+    
+    def update_count_for_words_sql(self, column_count):
+        return 'UPDATE %s SET %s=? WHERE %s' % (self.WORD_TABLE_NAME, self.COUNT_COL_NAME, self._make_column_names_and_placeholders(column_count)) 
 
-    def insert_row_for_words_sql(self):
-        columns = self._make_column_name_list() + ', ' + self.COUNT_COL_NAME
-        values  = ', '.join(['?'] * (self.column_count + 1))
+    def insert_row_for_words_sql(self, column_count):
+        columns = self._make_column_name_list(column_count) + ', ' + self.COUNT_COL_NAME
+        values  = ', '.join(['?'] * (column_count + 1))
 
-        return 'INSERT INTO %s (%s) VALUES (%s)' % (self.TABLE_NAME, columns, values) 
+        return 'INSERT INTO %s (%s) VALUES (%s)' % (self.WORD_TABLE_NAME, columns, values) 
 
-    def select_words_and_counts_sql(self):
-        last_word_col_name = self.WORD_COL_NAME_PREFIX + str(self.column_count)
+    def select_words_and_counts_sql(self, column_count):
+        last_word_col_name = self.WORD_COL_NAME_PREFIX + str(column_count)
 
-        return 'SELECT %s, %s FROM %s WHERE %s' % (last_word_col_name, self.COUNT_COL_NAME, self.TABLE_NAME, self._make_column_names_and_placeholders(self.column_count - 1))
+        return 'SELECT %s, %s FROM %s WHERE %s' % (last_word_col_name, self.COUNT_COL_NAME, self.WORD_TABLE_NAME, self._make_column_names_and_placeholders(column_count - 1))
 
     def delete_words_sql(self):
-        return 'DELETE FROM ' + self.TABLE_NAME
+        return 'DELETE FROM ' + self.WORD_TABLE_NAME
 

test/db_test.py

@@ -6,57 +6,64 @@ def setUp(self):
         self.conn = StubConn()
         self.sql = StubSql()
 
-    def test_correct_sql_run_when_object_created(self):
-        Db(3, self.conn, self.sql)
+    def test_correct_sql_run_when_setup_called(self):
+        Db(self.conn, self.sql).setup(3)
         execute_args = self.conn.stub_cursor.execute_args
-        self.assertEqual(len(execute_args), 2)
-        self.assertEqual(execute_args[0], ('create_table_sql',))
-        self.assertEqual(execute_args[1], ('create_index_sql',))
+        self.assertEqual(len(execute_args), 4)
+        self.assertEqual(execute_args[0], ('create_word_table_sql 3',))
+        self.assertEqual(execute_args[1], ('create_index_sql 3',))
+        self.assertEqual(execute_args[2], ('create_param_table_sql',))
+        self.assertEqual(execute_args[3], ('set_param_sql', ('depth', 3)))
 
     def test_error_when_add_word_count_wrong(self):
-        db = Db(3, self.conn, self.sql)
+        db = Db(self.conn, self.sql)
+        db.setup(3)
         self.assertRaises(ValueError, db.add_word, ['one','two'])
 
     def test_insert_row_when_add_new_word_list(self):
-        db = Db(3, self.conn, self.sql)
+        db = Db(self.conn, self.sql)
+        db.setup(3)
         word_list = ['one', 'two', 'three']
         db.add_word(word_list)
 
         execute_args = self.conn.stub_cursor.execute_args
-        self.assertEqual(len(execute_args), 4)
-        self.assertEqual(execute_args[2], ('select_count_for_words_sql', word_list))
-        self.assertEqual(execute_args[3], ('insert_row_for_words_sql', word_list + [1]))
+        self.assertEqual(len(execute_args), 6)
+        self.assertEqual(execute_args[4], ('select_count_for_words_sql 3', word_list))
+        self.assertEqual(execute_args[5], ('insert_row_for_words_sql 3', word_list + [1]))
 
     def test_update_row_when_add_repeated_word_list(self):
-        db = Db(3, self.conn, self.sql)
+        db = Db(self.conn, self.sql)
+        db.setup(3)
         row_count = 10
         word_list = ['one', 'two', 'three']
         self.conn.stub_cursor.fetchone_results.append([row_count])
 
         db.add_word(word_list)
 
         execute_args = self.conn.stub_cursor.execute_args
-        self.assertEqual(len(execute_args), 4)
-        self.assertEqual(execute_args[2], ('select_count_for_words_sql', word_list))
-        self.assertEqual(execute_args[3], ('update_count_for_words_sql', [row_count + 1] + word_list))
+        self.assertEqual(len(execute_args), 6)
+        self.assertEqual(execute_args[4], ('select_count_for_words_sql 3', word_list))
+        self.assertEqual(execute_args[5], ('update_count_for_words_sql 3', [row_count + 1] + word_list))
 
     def test_db_commit_performed_correctly(self):
-        db = Db(3, self.conn, self.sql)
+        db = Db(self.conn, self.sql)
+        db.setup(3)
         self.assertEqual(self.conn.commit_count, 0)
         db.commit()
         self.assertEqual(self.conn.commit_count, 1)
 
     def test_get_word_counts_works_correctly(self):
-        db = Db(3, self.conn, self.sql)
+        db = Db(self.conn, self.sql)
+        db.setup(3)
         word_list = ['i', 'like']
         self.conn.stub_cursor.execute_results = [[['dogs',  1], ['cats',  2], ['frogs', 3]]]
 
         word_counts = db.get_word_count(word_list)
 
         self.assertEqual(word_counts, {'dogs' : 1, 'cats' : 2, 'frogs' : 3})
         execute_args = self.conn.stub_cursor.execute_args
-        self.assertEqual(len(execute_args), 3)
-        self.assertEqual(execute_args[2], ('select_words_and_counts_sql', word_list))
+        self.assertEqual(len(execute_args), 5)
+        self.assertEqual(execute_args[4], ('select_words_and_counts_sql 3', word_list))
 
 class StubCursor:
     def __init__(self):
@@ -96,23 +103,32 @@ def cursor(self):
         return self.stub_cursor
 
 class StubSql:
-    def create_table_sql(self):
-        return 'create_table_sql'
+    def create_word_table_sql(self, column_count):
+        return 'create_word_table_sql' + ' ' + str(column_count)
+    
+    def create_index_sql(self, column_count):
+        return 'create_index_sql' + ' ' + str(column_count)
+    
+    def create_param_table_sql(self):
+        return 'create_param_table_sql'
+    
+    def set_param_sql(self):
+        return 'set_param_sql'
 
-    def create_index_sql(self):
-        return 'create_index_sql'
+    def get_param_sql(self):
+        return 'get_param_sql'
 
-    def select_count_for_words_sql(self):
-        return 'select_count_for_words_sql' 
+    def select_count_for_words_sql(self, column_count):
+        return 'select_count_for_words_sql' + ' ' + str(column_count)
 
-    def update_count_for_words_sql(self):
-        return 'update_count_for_words_sql' 
+    def update_count_for_words_sql(self, column_count):
+        return 'update_count_for_words_sql'  + ' ' + str(column_count)
 
-    def insert_row_for_words_sql(self):
-        return 'insert_row_for_words_sql' 
+    def insert_row_for_words_sql(self, column_count):
+        return 'insert_row_for_words_sql' + ' ' + str(column_count)
 
-    def select_words_and_counts_sql(self):
-        return 'select_words_and_counts_sql'
+    def select_words_and_counts_sql(self, column_count):
+        return 'select_words_and_counts_sql' + ' ' + str(column_count)
 
     def delete_words_sql(self):
         return 'delete_words_sql'            

test/gen_test.py

@@ -19,15 +19,19 @@ def test_generated_sequence_is_correct(self):
 
         self.rnd.vals = [1, 2, 2, 1, 4, 1, 1]
 
-        self.assertEqual(Generator('name', self.db, self.rnd).generate(2, ' '), 'the cat sat on the mat')
+        self.assertEqual(Generator('name', self.db, self.rnd).generate(' '), 'the cat sat on the mat')
         self.assertEqual(self.db.get_word_count_args, [['^'], ['the'], ['cat'], ['sat'], ['on'], ['the'], ['mat']])
         self.assertEqual(self.rnd.maxints, [3, 2, 2, 5, 4, 2, 1])
 
 class StubDb:
     def __init__(self):
         self.count_values = []
         self.get_word_count_args = []
+        self.depth = 2
 
+    def get_depth(self):
+        return self.depth
+    
     def get_word_count(self, word_list):   
         self.get_word_count_args.append(word_list)     
         return self.count_values.pop(0)