Add a bunch of functionality to !remind (#43).

earwigbot/commands/remind.py

@@ -20,43 +20,245 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from threading import Timer
+from itertools import chain
+import random
+from threading import Thread
 import time
 
 from earwigbot.commands import Command
 
+DISPLAY = ["display", "show", "list", "info"]
+CANCEL = ["cancel", "stop", "delete", "del"]
+SNOOZE = ["snooze", "delay", "reset", "adjust"]
+
 class Remind(Command):
     """Set a message to be repeated to you in a certain amount of time."""
     name = "remind"
-    commands = ["remind", "reminder"]
+    commands = ["remind", "reminder", "reminders", "snooze"]
 
-    def process(self, data):
-        if not data.args:
-            msg = "Please specify a time (in seconds) and a message in the following format: !remind <time> <msg>."
-            self.reply(data, msg)
-            return
+    @staticmethod
+    def _normalize(command):
+        """Convert a command name into its canonical form."""
+        if command in DISPLAY:
+            return "display"
+        if command in CANCEL:
+            return "cancel"
+        if command in SNOOZE:
+            return "snooze"
+
+    def _really_get_reminder_by_id(self, user, rid):
+        """Return the _Reminder object that corresponds to a particular ID.
+
+        Raises IndexError on failure.
+        """
+        if user not in self.reminders:
+            raise IndexError(rid)
+        return [robj for robj in self.reminders[user] if robj.id == rid][0]
+
+    def _get_reminder_by_id(self, user, rid, data):
+        """Return the _Reminder object that corresponds to a particular ID.
+
+        Sends an error message to the user on failure.
+        """
+        try:
+            return self._really_get_reminder_by_id(user, rid)
+        except IndexError:
+            msg = "Couldn't find a reminder for \x0302{0}\x0F with ID \x0303{1}\x0F."
+            self.reply(data, msg.format(user, rid))
+
+    def _get_new_id(self):
+        """Get a free ID for a new reminder."""
+        taken = set(robj.id for robj in chain(*self.reminders.values()))
+        num = random.choice(list(set(range(4096)) - taken))
+        return "R-{0:03X}".format(num)
 
+    def _create_reminder(self, data, user):
+        """Create a new reminder for the given user."""
         try:
             wait = int(data.args[0])
-        except ValueError:
-            msg = "The time must be given as an integer, in seconds."
-            self.reply(data, msg)
+            assert wait > 0
+        except (ValueError, AssertionError):
+            msg = "Invalid time \x02{0}\x0F. Time must be a positive integer, in seconds."
+            return self.reply(data, msg.format(data.args[0]))
+        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]
+        msg = "Set reminder \x0303{0}\x0F ({1})."
+        self.reply(data, msg.format(rid, wait, reminder.end_time))
+
+    def _display_reminder(self, data, reminder):
+        """Display a particular reminder's information."""
+        msg = 'Reminder \x0303{0}\x0F: {1} seconds ({2}): "{3}". !remind [cancel|snooze] {0}.'
+        msg = msg.format(reminder.id, reminder.wait, reminder.end_time,
+                         reminder.message)
+        self.reply(data, msg)
+
+    def _cancel_reminder(self, data, user, reminder):
+        """Cancel a pending reminder."""
+        reminder.stop()
+        self.reminders[user].remove(reminder)
+        if not self.reminders[user]:
+            del self.reminders[user]
+        msg = "Reminder \x0303{0}\x0F canceled."
+        self.reply(data, msg.format(reminder.id))
+
+    def _snooze_reminder(self, data, reminder, arg=None):
+        """Snooze a reminder to be re-triggered after a period of time."""
+        if arg:
+            try:
+                duration = int(data.args[arg])
+                assert duration > 0
+                reminder.wait = duration
+            except (IndexError, ValueError, AssertionError):
+                pass
+        reminder.start()
+        end_time = time.strftime("%b %d %H:%M:%S %Z", reminder.end)
+        msg = "Reminder \x0303{0}\x0F snoozed until {1}."
+        self.reply(data, msg.format(reminder.id, end_time))
+
+    def _handle_command(self, command, data, user, reminder, arg=None):
+        """Handle a reminder-processing subcommand."""
+        if command in DISPLAY:
+            self._display_reminder(data, reminder)
+        elif command in CANCEL:
+            self._cancel_reminder(data, user, reminder)
+        elif command in SNOOZE:
+            self._snooze_reminder(data, reminder, arg)
+        else:
+            msg = "Unknown action \x02{0}\x0F for reminder \x0303{1}\x0F."
+            self.reply(data, msg.format(command, reminder.id))
+
+    def _show_reminders(self, data, user):
+        """Show all of a user's current reminders."""
+        shorten = lambda s: (s[:37] + "..." if len(s) > 40 else s)
+        tmpl = '\x0303{0}\x0F ("{1}", {2})'
+        fmt = lambda robj: tmpl.format(robj.id, shorten(robj.message),
+                                       robj.end_time)
+
+        if user in self.reminders:
+            rlist = ", ".join(fmt(robj) for robj in self.reminders[user])
+            msg = "Your reminders: {1}.".format(rlist)
+        else:
+            msg = "You have no reminders. Set one with !remind <time> <msg>."
+        self.reply(data, msg)
+
+    def _process_snooze_command(self, data, user):
+        """Process the !snooze command."""
+        if not data.args:
+            if user not in self.reminders:
+                self.reply(data, "You have no reminders to snooze.")
+            elif len(self.reminders[user]) == 1:
+                self._snooze_reminder(data, self.reminders[user][0])
+            else:
+                msg = "You have {0} reminders. Snooze which one?"
+                self.reply(data, msg.format(len(self.reminders[user])))
+            return
+        reminder = self._get_reminder_by_id(user, data.args[0], data)
+        if reminder:
+            self._snooze_reminder(data, reminder, 1)
+
+    def setup(self):
+        self.reminders = {}
+
+    def process(self, data):
+        if data.command == "snooze":
+            return self._process_snooze_command(data, data.host)
+        if not data.args:
+            return self._show_reminders(data, data.host)
+
+        user = data.host
+        if len(data.args) == 1:
+            command = data.args[0]
+            if command in DISPLAY + CANCEL + SNOOZE:
+                if user not in self.reminders:
+                    msg = "You have no reminders to {0}."
+                    self.reply(data, msg.format(self._normalize(command)))
+                elif len(self.reminders[user]) == 1:
+                    reminder = self.reminders[user][0]
+                    self._handle_command(command, data, user, reminder)
+                else:
+                    msg = "You have {0} reminders. {1} which one?"
+                    num = len(self.reminders[user])
+                    command = self._normalize(command).capitalize()
+                    self.reply(data, msg.format(num, command))
+                return
+            reminder = self._get_reminder_by_id(user, data.args[0], data)
+            if reminder:
+                self._display_reminder(data, reminder)
             return
-        message = ' '.join(data.args[1:])
-        if not message:
-            msg = "What message do you want me to give you when time is up?"
-            self.reply(data, msg)
+
+        if data.args[0] in DISPLAY + CANCEL + SNOOZE:
+            reminder = self._get_reminder_by_id(user, data.args[1], data)
+            if reminder:
+                self._handle_command(data.args[0], data, user, reminder, 2)
             return
 
-        end = time.localtime(time.time() + wait)
-        end_time = time.strftime("%b %d %H:%M:%S", end)
-        end_time_with_timezone = time.strftime("%b %d %H:%M:%S %Z", end)
+        try:
+            reminder = self._really_get_reminder_by_id(user, data.args[0])
+        except IndexError:
+            return self._create_reminder(data, user)
 
-        msg = 'Set reminder for "{0}" in {1} seconds (ends {2}).'
-        msg = msg.format(message, wait, end_time_with_timezone)
-        self.reply(data, msg)
+        self._handle_command(data.args[1], data, user, reminder, 2)
 
-        t_reminder = Timer(wait, self.reply, args=(data, message))
-        t_reminder.name = "reminder " + end_time
-        t_reminder.daemon = True
-        t_reminder.start()
+
+class _Reminder(object):
+    """Represents a single reminder."""
+
+    def __init__(self, rid, user, wait, message, data, cmdobj):
+        self.id = rid
+        self.wait = wait
+        self.message = message
+        self.end = None
+
+        self._user = user
+        self._data = data
+        self._cmdobj = cmdobj
+        self._thread = None
+
+    def _callback(self):
+        """Internal callback function to be executed by the reminder thread."""
+        thread = self._thread
+        while time.time() < thread.end:
+            if thread.abort:
+                return
+            time.sleep(1)
+        time.sleep(60)
+        try:
+            self._cmdobj.reminders[self._user].remove(self)
+            if not self._cmdobj.reminders[self._user]:
+                del self._cmdobj.reminders[self._user]
+        except (KeyError, ValueError):  # Already canceled by the user
+            pass
+
+    @property
+    def end_time(self):
+        """Return a string representing the end time of a reminder."""
+        if self.end >= time.time():
+            end_time = time.strftime("%b %d %H:%M:%S %Z", self.end)
+            return "ends {0}".format(end_time)
+        return "expired"
+
+    def start(self):
+        """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.daemon = True
+        self._thread.abort = False
+        self._thread.start()
+
+    def stop(self):
+        """Stop a currently running reminder."""
+        if not self._thread:
+            return
+        self._thread.abort = True
+        self._thread = None