nxy/bot/plugins/timer.py

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

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

from . import DatabasePlugin


@irc3.plugin
class Timer(DatabasePlugin):
    requires = ['irc3.plugins.command',
                'bot.plugins.storage']

    def __init__(self, bot: irc3.IrcBot):
        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:
            self.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(self.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')
        self.cur.execute('''
        SELECT
          *
        FROM
          timers
        WHERE
          ends_at >= now()
          AND ends_at < now() + INTERVAL '1h'
        ''')

        for timer in self.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'])
        self.cur.execute('''
        DELETE FROM
          timers
        WHERE
          id = %s
        ''', [timer['id']])
        self.con.commit()