Make reminders persistent.

earwigbot/commands/remind.py

@@ -21,13 +21,15 @@
 # SOFTWARE.
 
 import ast
+from contextlib import contextmanager
 from itertools import chain
 import operator
 import random
-from threading import Thread
+from threading import RLock, Thread
 import time
 
 from earwigbot.commands import Command
+from earwigbot.irc import Data
 
 DISPLAY = ["display", "show", "list", "info", "details"]
 CANCEL = ["cancel", "stop", "delete", "del", "stop", "unremind", "forget",
@@ -60,6 +62,7 @@ class Remind(Command):
             ast.Pow: operator.pow
         }
         time_units = {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}
+
         def _evaluate(node):
             """Convert an AST node into a real number or raise an exception."""
             if isinstance(node, ast.Num):
@@ -76,6 +79,7 @@ class Remind(Command):
             factor, arg = time_units[arg[-1]], arg[:-1]
         else:
             factor = 1
+
         try:
             parsed = int(_evaluate(ast.parse(arg, mode="eval").body) * factor)
         except (SyntaxError, KeyError):
@@ -84,6 +88,12 @@ class Remind(Command):
             raise ValueError(parsed)
         return parsed
 
+    @contextmanager
+    def _db(self):
+        """Return a threadsafe context manager for the permissions database."""
+        with self._db_lock:
+            yield self.config.irc["permissions"]
+
     def _really_get_reminder_by_id(self, user, rid):
         """Return the _Reminder object that corresponds to a particular ID.
 
@@ -111,6 +121,14 @@ class Remind(Command):
         num = random.choice(list(set(range(4096)) - taken))
         return "R{0:03X}".format(num)
 
+    def _start_reminder(self, reminder, user):
+        """Start the given reminder object for the given user."""
+        reminder.start()
+        if user in self.reminders:
+            self.reminders[user].append(reminder)
+        else:
+            self.reminders[user] = [reminder]
+
     def _create_reminder(self, data, user):
         """Create a new reminder for the given user."""
         try:
@@ -118,22 +136,22 @@ class Remind(Command):
         except ValueError:
             msg = "Invalid time \x02{0}\x0F. Time must be a positive integer, in seconds."
             return self.reply(data, msg.format(data.args[0]))
+
         if wait > 1000 * 365 * 24 * 60 * 60:
             # Hard to think of a good upper limit, but 1000 years works.
             msg = "Given time \x02{0}\x0F is too large. Keep it reasonable."
             return self.reply(data, msg.format(data.args[0]))
+
+        end = time.time() + wait
         message = " ".join(data.args[1:])
         try:
             rid = self._get_new_id()
         except IndexError:
             msg = "Couldn't set a new reminder: no free IDs available."
             return self.reply(data, msg)
-        reminder = _Reminder(rid, user, wait, message, data, self)
-        reminder.start()
-        if user in self.reminders:
-            self.reminders[user].append(reminder)
-        else:
-            self.reminders[user] = [reminder]
+
+        reminder = _Reminder(rid, user, wait, end, message, data, self)
+        self._start_reminder(reminder, user)
         msg = "Set reminder \x0303{0}\x0F ({1})."
         self.reply(data, msg.format(rid, reminder.end_time))
 
@@ -162,11 +180,30 @@ class Remind(Command):
                 reminder.wait = duration
             except (IndexError, ValueError):
                 pass
+
+        reminder.end = time.time() + reminder.wait
         reminder.start()
         end = time.strftime("%b %d %H:%M:%S %Z", time.localtime(reminder.end))
         msg = "Reminder \x0303{0}\x0F {1} until {2}."
         self.reply(data, msg.format(reminder.id, verb, end))
 
+    def _load_reminders(self):
+        """Load previously made reminders from the database."""
+        with self._db() as permdb:
+            try:
+                database = permdb.get_attr("command:remind", "data")
+            except KeyError:
+                return
+            permdb.set_attr("command:remind", "data", "[]")
+
+        for item in ast.literal_eval(database):
+            rid, user, wait, end, message, data = item
+            if end < time.time():
+                continue
+            data = Data.unserialize(data)
+            reminder = _Reminder(rid, user, wait, end, message, data, self)
+            self._start_reminder(reminder, user)
+
     def _handle_command(self, command, data, user, reminder, arg=None):
         """Handle a reminder-processing subcommand."""
         if command in DISPLAY:
@@ -190,7 +227,8 @@ class Remind(Command):
             rlist = ", ".join(fmt(robj) for robj in self.reminders[user])
             msg = "Your reminders: {0}.".format(rlist)
         else:
-            msg = "You have no reminders. Set one with \x0306!remind [time] [message]\x0F. See also: \x0306!remind help\x0F."
+            msg = ("You have no reminders. Set one with \x0306!remind [time] "
+                   "[message]\x0F. See also: \x0306!remind help\x0F.")
         self.reply(data, msg)
 
     def _process_snooze_command(self, data, user):
@@ -239,6 +277,8 @@ class Remind(Command):
 
     def setup(self):
         self.reminders = {}
+        self._db_lock = RLock()
+        self._load_reminders()
 
     def process(self, data):
         if data.command == "snooze":
@@ -284,15 +324,42 @@ class Remind(Command):
 
         self._handle_command(data.args[1], data, user, reminder, 2)
 
+    def unload(self):
+        for reminder in chain(*self.reminders.values()):
+            reminder.stop(delete=False)
+
+    def store_reminder(self, reminder):
+        """Store a serialized reminder into the database."""
+        with self._db() as permdb:
+            try:
+                dump = permdb.get_attr("command:remind", "data")
+            except KeyError:
+                dump = "[]"
+
+            database = ast.literal_eval(dump)
+            database.append(reminder)
+            permdb.set_attr("command:remind", "data", str(database))
+
+    def unstore_reminder(self, rid):
+        """Remove a reminder from the database by ID."""
+        with self._db() as permdb:
+            try:
+                dump = permdb.get_attr("command:remind", "data")
+            except KeyError:
+                dump = "[]"
+
+            database = ast.literal_eval(dump)
+            database = [item for item in database if item[0] != rid]
+            permdb.set_attr("command:remind", "data", str(database))
 
 class _Reminder(object):
     """Represents a single reminder."""
 
-    def __init__(self, rid, user, wait, message, data, cmdobj):
+    def __init__(self, rid, user, wait, end, message, data, cmdobj):
         self.id = rid
         self.wait = wait
+        self.end = end
         self.message = message
-        self.end = None
 
         self._user = user
         self._data = data
@@ -307,6 +374,7 @@ class _Reminder(object):
             if thread.abort:
                 return
         self._cmdobj.reply(self._data, self.message)
+        self._delete()
         for i in xrange(60):
             time.sleep(1)
             if thread.abort:
@@ -318,6 +386,16 @@ class _Reminder(object):
         except (KeyError, ValueError):  # Already canceled by the user
             pass
 
+    def _save(self):
+        """Save this reminder to the database."""
+        data = self._data.serialize()
+        item = (self.id, self._user, self.wait, self.end, self.message, data)
+        self._cmdobj.store_reminder(item)
+
+    def _delete(self):
+        """Remove this reminder from the database."""
+        self._cmdobj.unstore_reminder(self.id)
+
     @property
     def end_time(self):
         """Return a string representing the end time of a reminder."""
@@ -330,14 +408,17 @@ class _Reminder(object):
         """Start the reminder timer thread. Stops it if already running."""
         self.stop()
         self._thread = Thread(target=self._callback, name="remind-" + self.id)
-        self._thread.end = self.end = time.time() + self.wait
+        self._thread.end = self.end
         self._thread.daemon = True
         self._thread.abort = False
         self._thread.start()
+        self._save()
 
-    def stop(self):
+    def stop(self, delete=True):
         """Stop a currently running reminder."""
         if not self._thread:
             return
+        if delete:
+            self._delete()
         self._thread.abort = True
         self._thread = None