nxy/bot/timer.py

# -*- coding: utf-8 -*-
import re
import asyncio
from datetime import datetime

from aiocron import crontab
from docopt import Dict
from irc3.plugins.command import command
from irc3.utils import IrcString
from psycopg2 import Error
from psycopg2.extras import DictRow

from . import DatabasePlugin, Bot


class Timer(DatabasePlugin):
    def __init__(self, bot: Bot):
        super().__init__(bot)
        self.timers = set()
        self.set_timers()
        crontab('0 * * * *', func=self.set_timers)

    @command
    def timer(self, mask: IrcString, target: IrcString, args: Dict):
        """Sets a timer, delay can be: s, m, h, d, w, mon, y

        %%timer <delay> <message>...
        """
        delay = args['<delay>']
        message = ' '.join(args['<message>'])

        # TODO: allow precise delays
        if not re.match(r'\d+[smhdwy]|mon', delay):
            return 'Invalid timer delay: {}'.format(delay)

        try:
            with self.con.cursor() as cur:
                cur.execute('''
                INSERT INTO
                  timers (mask, target, message, delay, ends_at)
                VALUES
                  (%s, %s, %s, %s, now() + INTERVAL %s)
                RETURNING
                  *
                ''', [mask, target, message, delay, delay])
            self.con.commit()

            asyncio.ensure_future(self.exec_timer(cur.fetchone()))

            self.bot.notice(mask.nick, 'Timer in {delay} set: {message}'.format(delay=delay, message=message))
        except Error as ex:
            self.log.error(ex)
            self.con.rollback()

    def set_timers(self):
        """Function which queries all timers in the next hour and schedules them."""
        self.log.debug('Fetching timers')
        with self.con.cursor() as cur:
            cur.execute('''
            SELECT
              *
            FROM
              timers
            WHERE
              ends_at >= now()
              AND ends_at < now() + INTERVAL '1h'
            ''')

            for timer in cur.fetchall():
                asyncio.ensure_future(self.exec_timer(timer))

    async def exec_timer(self, timer: DictRow):
        """Sets the actual timer (sleeps until it fires), sends the reminder and deletes the timer from database."""
        if timer['id'] in self.timers:
            return

        self.timers.add(timer['id'])
        seconds = (timer['ends_at'] - datetime.now()).total_seconds()
        if seconds > 0.0:
            await asyncio.sleep(seconds)

        self.bot.privmsg(timer['target'], '\x02[Timer]\x02 {nick}: {message} ({delay})'.format(
            message=timer['message'],
            nick=IrcString(timer['mask']).nick,
            delay=timer['delay'],
        ))

        self.timers.remove(timer['id'])
        with self.con.cursor() as cur:
            cur.execute('''
            DELETE FROM
              timers
            WHERE
              id = %s
            ''', [timer['id']])
        self.con.commit()