Initial commit

.gitignore

+*.mbp

README.md

+# reminder
+A [maubot](https://github.com/maubot/maubot) to remind you about things.

base-config.yaml

+# Base command without the prefix (!).
+base_command: remind

maubot.yaml

+maubot: 0.1.0
+id: xyz.maubot.reminder
+version: 0.1.0
+license: AGPL-3.0-or-later
+modules:
+- reminder
+main_class: ReminderBot
+extra_files:
+- base-config.yaml
+dependencies:
+- dateparser>=0.7.1,<0.8
+- pytz
+database: true

reminder/__init__.py

+from .bot import ReminderBot

reminder/bot.py

+# reminder - A maubot plugin to remind you about things.
+# Copyright (C) 2019 Tulir Asokan
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from typing import Type
+from datetime import datetime, timedelta
+from html import escape
+import asyncio
+
+import pytz
+
+from mautrix.util.config import BaseProxyConfig
+from maubot import Plugin, MessageEvent
+from maubot.handlers import command
+
+from .db import ReminderDatabase
+from .util import Config, ReminderInfo, DateArgument
+
+
+class ReminderBot(Plugin):
+    db: ReminderDatabase
+    reminder_loop_task: asyncio.Future
+
+    @classmethod
+    def get_config_class(cls) -> Type[BaseProxyConfig]:
+        return Config
+
+    async def start(self) -> None:
+        await super().start()
+        self.config.load_and_update()
+        self.db = ReminderDatabase(self.database)
+        self.reminder_loop_task = asyncio.ensure_future(self.reminder_loop(), loop=self.loop)
+
+    async def stop(self) -> None:
+        await super().stop()
+        self.reminder_loop_task.cancel()
+
+    async def reminder_loop(self) -> None:
+        try:
+            self.log.debug("Reminder loop started")
+            while True:
+                now = datetime.now(tz=pytz.UTC)
+                next_minute = (now + timedelta(minutes=1)).replace(second=0, microsecond=0)
+                await asyncio.sleep((next_minute - now).total_seconds())
+                await self.schedule_nearby_reminders(next_minute)
+        except asyncio.CancelledError:
+            self.log.debug("Reminder loop stopped")
+        except Exception:
+            self.log.exception("Exception in reminder loop")
+
+    async def schedule_nearby_reminders(self, now: datetime) -> None:
+        until = now + timedelta(minutes=1)
+        for reminder in self.db.all_in_range(now, until):
+            asyncio.ensure_future(self.send_reminder(reminder), loop=self.loop)
+
+    async def send_reminder(self, reminder: ReminderInfo) -> None:
+        try:
+            await self._send_reminder(reminder)
+        except Exception:
+            self.log.exception("Failed to send reminder")
+
+    async def _send_reminder(self, reminder: ReminderInfo) -> None:
+        if len(reminder.users) == 0:
+            self.log.debug(f"Cancelling reminder {reminder}, no users left to remind")
+            return
+        wait = (reminder.date - datetime.now(tz=pytz.UTC)).total_seconds()
+        if wait > 0:
+            self.log.debug(f"Waiting {wait} seconds to send {reminder}")
+            await asyncio.sleep(wait)
+        else:
+            self.log.debug(f"Sending {reminder} immediately")
+        users = " ".join(reminder.users)
+        users_html = " ".join(f"<a href='https://matrix.to/#/{user_id}'>{user_id}</a>"
+                              for user_id in reminder.users)
+        text = f"{users}: {reminder.message}"
+        html = f"{users_html}: {escape(reminder.message)}"
+        await self.client.send_text(reminder.room_id, text=text, html=html)
+
+    @command.new(name=lambda self: self.config["base_command"],
+                 help="Create a reminder", require_subcommand=False, arg_fallthrough=False)
+    @DateArgument("date", required=True)
+    @command.argument("message", pass_raw=True, required=True)
+    async def remind(self, evt: MessageEvent, date: datetime, message: str) -> None:
+        date = date.replace(microsecond=0)
+        if date < datetime.now(tz=pytz.UTC):
+            await evt.reply(f"Sorry, {date} is in the past and I don't have a time machine :(")
+            return
+        rem = ReminderInfo(date=date, room_id=evt.room_id, message=message, users=[evt.sender])
+        self.db.insert(rem)
+        await evt.reply(f"Reminder #{rem.id} will remind you for {rem.message} at {rem.date}.")
+        now = datetime.now(tz=pytz.UTC)
+        if (date - now).total_seconds() < 60 and now.minute == date.minute:
+            self.log.debug(f"Reminder {rem} is in less than a minute, scheduling now...")
+            asyncio.ensure_future(self.send_reminder(rem), loop=self.loop)
+
+    @remind.subcommand("list", help="List your reminders")
+    async def list(self, evt: MessageEvent) -> None:
+        reminders_str = "\n".join(f"* {reminder.message} at {reminder.date}"
+                                  for reminder in self.db.all_for_user(evt.sender))
+        if len(reminders_str) == 0:
+            await evt.reply("You have no upcoming reminders :(")
+        else:
+            await evt.reply(f"List of reminders:\n\n{reminders_str}")
+
+    @remind.subcommand("cancel", help="Cancel a reminder", aliases=("delete", "remove", "rm"))
+    @command.argument("id", parser=lambda val: int(val) if val else None)
+    async def cancel(self, evt: MessageEvent, id: int) -> None:
+        reminder = self.db.get(id)
+        self.db.remove_user(reminder, evt.sender)
+        await evt.reply(f"Reminder #{reminder.id}: {reminder.message} at {reminder.date} cancelled")

reminder/db.py

+# reminder - A maubot plugin to remind you about things.
+# Copyright (C) 2019 Tulir Asokan
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from typing import Optional, Iterator, Dict
+from datetime import datetime
+
+import pytz
+from sqlalchemy import (Column, String, Integer, Text, DateTime, ForeignKey, Table, MetaData,
+                        select, and_)
+from sqlalchemy.engine.base import Engine
+
+from mautrix.types import UserID
+
+from .util import ReminderInfo
+
+
+class ReminderDatabase:
+    reminder: Table
+    reminder_target: Table
+    timezone: Table
+    tz_cache: Dict[UserID, pytz.timezone]
+    db: Engine
+
+    def __init__(self, db: Engine) -> None:
+        self.db = db
+
+        meta = MetaData()
+        meta.bind = db
+
+        self.reminder = Table("reminder", meta,
+                              Column("id", Integer, primary_key=True, autoincrement=True),
+                              Column("date", DateTime, nullable=False),
+                              Column("room_id", String(255), nullable=False),
+                              Column("message", Text, nullable=False))
+        self.reminder_target = Table("reminder_target", meta,
+                                     Column("reminder_id", Integer,
+                                            ForeignKey("reminder.id", ondelete="CASCADE"),
+                                            primary_key=True),
+                                     Column("user_id", String(255), primary_key=True))
+        self.timezone = Table("timezone", meta,
+                              Column("user_id", String(255), primary_key=True),
+                              Column("timezone", String(255), primary_key=True))
+
+        meta.create_all()
+
+    def set_timezone(self, user_id: UserID, tz: pytz.timezone) -> None:
+        self.db.execute(self.timezone.insert().values(user_id=user_id, timezone=tz.zone))
+        self.tz_cache[user_id] = tz
+
+    def get_timezone(self, user_id: UserID) -> Optional[pytz.timezone]:
+        try:
+            return self.tz_cache[user_id]
+        except KeyError:
+            rows = self.db.execute(self.timezone.select().where(self.timezone.c.user_id == user_id))
+            try:
+                self.tz_cache[user_id] = pytz.timezone(next(rows)[0])
+            except (pytz.UnknownTimeZoneError, StopIteration, IndexError):
+                self.tz_cache[user_id] = None
+            return self.tz_cache[user_id]
+
+    def all_for_user(self, user_id: UserID) -> Iterator[ReminderInfo]:
+        rows = self.db.execute(select([self.reminder]).where(
+            and_(self.reminder.c.id == self.reminder_target.c.reminder_id,
+                 self.reminder_target.c.user_id == user_id,
+                 self.reminder.c.date > datetime.now(tz=pytz.UTC))))
+        for row in rows:
+            yield ReminderInfo(id=row[0], date=row[1].replace(tzinfo=pytz.UTC), room_id=row[2],
+                               message=row[3], users=[user_id])
+
+    def get(self, id: int) -> Optional[ReminderInfo]:
+        rows = self.db.execute(select([self.reminder, self.reminder_target.c.user_id]).where(
+            and_(self.reminder.c.id == id,
+                 self.reminder_target.c.reminder_id == self.reminder.c.id)))
+        try:
+            first_row = next(rows)
+        except StopIteration:
+            return None
+        info = ReminderInfo(id=first_row[0], date=first_row[1].replace(tzinfo=pytz.UTC),
+                            room_id=first_row[2], message=first_row[3], users=[first_row[4]])
+        for row in rows:
+            info.users.append(row[4])
+        return info
+
+    def _get_many(self, whereclause) -> Iterator[ReminderInfo]:
+        rows = self.db.execute(select([self.reminder, self.reminder_target.c.user_id])
+                               .where(whereclause)
+                               .order_by(self.reminder.c.id, self.reminder.c.date))
+        building_reminder = None
+        for row in rows:
+            if building_reminder is not None:
+                if building_reminder.id == row[0]:
+                    building_reminder.users.append(row[4])
+                    continue
+                yield building_reminder
+            building_reminder = ReminderInfo(id=row[0], date=row[1].replace(tzinfo=pytz.UTC),
+                                             room_id=row[2], message=row[3], users=[row[4]])
+        if building_reminder is not None:
+            yield building_reminder
+
+    def all(self) -> Iterator[ReminderInfo]:
+        yield from self._get_many(self.reminder_target.c.reminder_id == self.reminder.c.id)
+
+    def all_in_range(self, after: datetime, before: datetime) -> Iterator[ReminderInfo]:
+        yield from self._get_many(and_(self.reminder_target.c.reminder_id == self.reminder.c.id,
+                                       after <= self.reminder.c.date,
+                                       self.reminder.c.date < before))
+
+    def insert(self, reminder: ReminderInfo) -> None:
+        with self.db.begin() as tx:
+            res = tx.execute(self.reminder.insert()
+                             .values(date=reminder.date, room_id=reminder.room_id,
+                                     message=reminder.message))
+            reminder.id = res.inserted_primary_key[0]
+            print([{"reminder_id": reminder.id, "user_id": user_id}
+                   for user_id in reminder.users])
+            tx.execute(self.reminder_target.insert(),
+                       [{"reminder_id": reminder.id, "user_id": user_id}
+                        for user_id in reminder.users])
+
+    def add_user(self, reminder: ReminderInfo, user_id: UserID) -> None:
+        if user_id in reminder.users:
+            return
+        self.db.execute(self.reminder_target.insert()
+                        .values(reminder_id=reminder.id, user_id=user_id))
+        reminder.users.append(user_id)
+
+    def remove_user(self, reminder: ReminderInfo, user_id: UserID) -> None:
+        if user_id not in reminder.users:
+            return
+        self.db.execute(self.reminder_target.delete().where(
+            and_(self.reminder_target.c.reminder_id == reminder.id,
+                 self.reminder_target.c.user_id == user_id)))
+        reminder.users.remove(user_id)

reminder/util.py

+# reminder - A maubot plugin to remind you about things.
+# Copyright (C) 2019 Tulir Asokan
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from typing import Optional, List, Tuple, TYPE_CHECKING
+from collections import deque
+from datetime import datetime
+from attr import dataclass
+
+import pytz
+import dateparser
+
+from mautrix.types import UserID, RoomID
+from mautrix.util.config import BaseProxyConfig, ConfigUpdateHelper
+from maubot import MessageEvent
+from maubot.handlers.command import Argument
+
+if TYPE_CHECKING:
+    from .bot import ReminderBot
+
+
+class Config(BaseProxyConfig):
+    def do_update(self, helper: ConfigUpdateHelper) -> None:
+        helper.copy("base_command")
+
+
+class DateArgument(Argument):
+    def __init__(self, name: str, label: str = None, *, required: bool = False):
+        super().__init__(name, label=label, required=required, pass_raw=True)
+
+    def match(self, val: str, evt: MessageEvent = None, instance: 'ReminderBot' = None
+              ) -> Tuple[str, Optional[datetime]]:
+        parser_settings = {
+            "PREFER_DATES_FROM": "future",
+            "TIMEZONE": "UTC",
+            "TO_TIMEZONE": "UTC",
+            "RETURN_AS_TIMEZONE_AWARE": True,
+        }
+
+        if instance:
+            tz = instance.db.get_timezone(evt.sender)
+            if tz:
+                parser_settings["TIMEZONE"] = tz.zone
+
+        time = None
+        parts = [part for part in val.split(" ") if len(part) > 0] + [""]
+        rem = deque()
+        while not time:
+            if len(parts) == 0:
+                return val, None
+            rem.appendleft(parts.pop())
+            time = dateparser.parse(" ".join(parts), settings=parser_settings)
+        if time < datetime.now(tz=pytz.UTC) and parts[0] != "in":
+            parts.insert(0, "in")
+            in_time = dateparser.parse(" ".join(parts), settings=parser_settings)
+            if in_time:
+                time = in_time
+        return " ".join(rem), time
+
+
+@dataclass
+class ReminderInfo:
+    id: int = None
+    date: datetime = None
+    message: str = None
+    room_id: RoomID = None
+    users: List[UserID] = None