Remember idle status and avatar in sqlite.

diff --git a/bot.py b/bot.py
index b148a49..5ed24b0 100755 (executable)
--- a/bot.py
+++ b/bot.py
@@ -17,6 +17,7 @@ import re
 import sys
 import time
 from PIL import Image
+from db import DBConnection
 
 dryrun = os.getenv('DRYRUN') is not None
 debug = os.getenv('DEBUG') is not None
@@ -26,6 +27,8 @@ if sys.version_info >= (3, 0):
 else:
   import urllib2
 
+db = DBConnection()
+
 token_file = 'TOKEN'
 if len(sys.argv) > 1:
   token_file = sys.argv[1]
@@ -150,18 +153,20 @@ def set_idle(idle):
 
 @asyncio.coroutine
 def wake_up(force = False):
+  last_spoke.value = time.time()
   if force or status() != discord.Status.online:
     log.info('Waking up...')
-    yield from set_avatar('avatar_online.png')
+    yield from set_avatar(db.ONLINE_AVATAR)
     yield from set_idle(False)
-  last_spoke.value = time.time()
+    db.set_state(client, avatar = db.ONLINE_AVATAR, idle = False, last_spoke = last_spoke.value)
 
 @asyncio.coroutine
 def zzz():
   if status() != discord.Status.idle:
     log.info('Sleeping...')
-    yield from set_avatar('avatar_idle.png')
+    yield from set_avatar(db.IDLE_AVATAR)
     yield from set_idle(True)
+    db.set_state(client, avatar = db.IDLE_AVATAR, idle = True)
 
 @asyncio.coroutine
 def maybe_say(channel, text, *, probability = 0.2, wake = True):
@@ -755,11 +760,22 @@ def play_greeting(greeting):
 @client.event
 @asyncio.coroutine
 def on_ready():
+  state = db.get_state(client)
+  if state is not None:
+    if state['idle']:
+      log.info('Restoring idle status.')
+      yield from set_idle(True)
+      if state['avatar'] != db.IDLE_AVATAR:
+        log.info('Restoring idle avatar.')
+        yield from set_avatar(db.IDLE_AVATAR)
+    elif state['avatar'] != db.ONLINE_AVATAR:
+      log.info('Restoring online avatar.')
+      yield from set_avatar(db.ONLINE_AVATAR)
+    last_spoke.value = state['last_spoke']
   log.info('Logged in as {}#{}'.format(client.user.name, client.user.id))
   feeds = get_bot_variable('feeds', [])
   orders = get_bot_variable('orders', [])
   log.info('Admins are: {}'.format(admins))
-  yield from wake_up(True)
   channels = channels_to_greet()
   log.info('Greeting in {}'.format([channel.name for channel in channels]))
   for channel in channels:

diff --git a/db.py b/db.py
new file mode 100644 (file)
index 0000000..64515e2
--- /dev/null
+++ b/db.py
@@ -0,0 +1,51 @@
+import discord
+import logging
+import sqlite3
+
+log = logging.getLogger('db')
+
+class DBConnection(object):
+  def __init__(self, filename = 'bot.sqlite'):
+    self.ONLINE_AVATAR = 'avatar_online.png'
+    self.IDLE_AVATAR = 'avatar_idle.png'
+    self.open_db(filename)
+    self.create_tables()
+
+  def open_db(self, filename):
+    self.dbh = sqlite3.connect(filename)
+    self.dbh.row_factory = sqlite3.Row
+
+  def create_tables(self):
+    cursor = self.query('create table if not exists state (bot_id varchar(32) not null, avatar varchar(128), idle boolean not null default false, last_spoke datetime)')
+    self.dbh.commit()
+
+  def close(self):
+    self.dbh.close()
+
+  def query(self, sql, params = ()):
+    log.debug(sql, params)
+    cursor = self.dbh.cursor()
+    cursor.execute(sql, params)
+    return cursor
+
+  def get_state(self, client):
+    cursor = self.dbh.cursor()
+    self.query('select avatar, idle, last_spoke from state where bot_id=?', [client.user.id])
+    row = cursor.fetchone()
+    if row is None:
+      return None
+    return dict(row)
+
+  def set_state(self, client, **args):
+    values = []
+    params = []
+    for k, v in args.items():
+      values.append(k)
+      params.append(v)
+    params.append(client.user.id)
+    sql = 'update state set {} where bot_id=?'.format(', '.join(['{}=?'.format(v) for v in values]))
+    cursor = self.query(sql, params)
+    if not cursor.rowcount:
+      cursor = self.query('insert into state (bot_id) values (?)', [client.user.id])
+      cursor = self.query(sql, params)
+    self.dbh.commit()