ameliabot/ameliabot/quote.py

69 lines
2.2 KiB
Python

from datetime import datetime, timezone
import random
import sqlite3
from ameliabot.logger import logging
class Quote:
def __init__(self):
self.__init_table()
self.num_quotes = self._get_num_quotes()
logging.info("Quote subsystem online")
def insert(self, owner, submitter, text):
conn = self.__connect()
text = text.replace("'", "''")
conn.execute('''
INSERT INTO quotes (submitter, text, timestamp)
VALUES ('{}', '{}', {})'''.format(
submitter, text, datetime.now().replace(tzinfo=timezone.utc)))
self.num_quotes += 1
logging.debug("Quote number %s inserted" % self.num_quotes)
def get(self, arg=None):
conn = self.__connect()
if arg:
ret = "SELECT id, text, timestamp FROM quotes WHERE "
try:
if int(arg) > self.num_quotes:
return "No quote matching that number."
ret += "id = %s" % arg
except ValueError:
ret += "text like '%{}%'".format(arg.lower())
quote = conn.execute(ret)
return self._format(list(quote.fetchone()))
else:
return self.get(random.randint(0, self.num_quotes))
def _get_num_quotes(self):
conn = self.__connect()
try:
num = conn.execute("SELECT COUNT(id) FROM quotes")
except sqlite3.OperationalError:
return 0
return int(num.fetchone()[0])
def _format(self, quote):
num, text, timestamp = quote
timestamp = datetime.strptime(timestamp, '%Y-%m-%d %H:%M:%S.%f%z')
return "{}. {}, {}".format(
num, text, datetime.strftime(timestamp, '%Y'))
def __connect(self):
return sqlite3.connect("data/quote.db")
def __init_table(self):
try:
conn = self.__connect()
except sqlite3.OperationalError:
import os
os.makedirs("data")
self.__init_table()
conn.execute('''
CREATE TABLE IF NOT EXISTS quotes (
id INTEGER PRIMARY KEY, submitter TEXT,
text TEXT, timestamp TEXT
)''')