import sqlite3

'''

create table word(word, doctype, count);

create table doctype_count(doctype, count);

create index i1 on word(word, doctype);

delete from word;

update ad_count set count = 0;

'''

class Db:

def __init__(self):

self.conn = sqlite3.connect('./bayes.db')

def reset(self):

c = self.conn.cursor()

try:

c.execute('delete from word')

c.execute('delete from doctype_count')

finally:

c.close()

self.conn.commit()

def update_word_count(self, c, doctype, word, num_to_add_to_count):

c.execute('select count from word where doctype=? and word=?', (doctype, word))

r = c.fetchone()

if r:

c.execute('update word set count=? where doctype=? and word=?', (r[0] + num_to_add_to_count, doctype, word))

else:

c.execute('insert into word (doctype, word, count) values (?,?,?)', (doctype, word, num_to_add_to_count))

def update_word_counts(self, d, doctype):

c = self.conn.cursor()

try:

for word, count in d.items():

self.update_word_count(c, doctype, word, count)

finally:

c.close()

self.conn.commit()

def get_doctype_counts(self):

counts = {}

c = self.conn.cursor()

try:

for row in c.execute('select doctype, count from doctype_count'):

counts[row[0]] = row[1]

return counts

finally:

c.close()

self.conn.commit()

def get_word_count(self, doctype, word):

c = self.conn.cursor()

try:

c.execute('select count from word where doctype=? and word=?', (doctype, word))

r = c.fetchone()

if r:

return r[0]

else:

return 0

finally:

c.close()

self.conn.commit()

def get_words_count(self, doctype):

c = self.conn.cursor()

try:

c.execute('select sum(count) from word where doctype=?', (doctype, ))

r = c.fetchone()

if r:

return r[0]

else:

return 0

finally:

c.close()

self.conn.commit()

def update_doctype_count(self, num_new_ads, doctype):

c = self.conn.cursor()

try:

counts = self.get_doctype_counts()

if counts.has_key(doctype):

current_count = counts[doctype]

else:

current_count = 0

if current_count:

c.execute('update doctype_count set count=? where doctype=?', (current_count + num_new_ads, doctype))

else:

c.execute('insert into doctype_count (doctype, count) values (?, ?)', (doctype, num_new_ads))

finally:

c.close()

self.conn.commit()