nxy/bot/plugins/timer.py

# -*- coding: utf-8 -*-
import asyncio
from datetime import datetime, timedelta

import irc3
from docopt import Dict as DocOptDict
from irc3.plugins.command import command
from irc3.utils import IrcString
from psycopg2 import Error

from . import DatabasePlugin
from ..utils import time_delta


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

    def __init__(self, bot: irc3.IrcBot):
        super().__init__(bot)

        # Fetch timers from database
        self.cur.execute('''
        select
          id, mask, target, message, delay, ends_at
        from
          timers
        ''')

        # Recreate timers
        for res in self.cur.fetchall():
            self.start_timer(IrcString(res['mask']),
                             res['target'],
                             res['message'],
                             res['delay'],
                             res['ends_at'] - datetime.now(),
                             res['id'])

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

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

        if not delta:
            self.bot.privmsg(target, 'Invalid timer delay')
        else:
            message = ' '.join(args['<message>'])
            values = [mask, target, message, delay]

            try:
                # Insert into database (add now + delta)
                self.cur.execute('''
                insert into
                  timers (mask, target, message, delay, ends_at)
                values
                  (%s, %s, %s, %s, %s)
                returning id
                ''', values + [datetime.now() + delta])
                self.con.commit()

                # Add delta and id from inserted and start timer
                values.extend([delta, self.cur.fetchone()['id']])
                self.start_timer(*values)

                # Send notice to user that timer has been set
                self.bot.notice(mask.nick, 'Timer in {delay} set: {message}'
                                .format(delay=delay, message=message))
            except Error:
                # Rollback transaction on error
                self.con.rollback()

    def start_timer(self, mask: IrcString, target: IrcString, message: str,
                    delay: str, delta: timedelta, row_id: int):
        """Async function, sleeps for `delay` seconds and sends notification"""

        async def callback():
            # Sleep if necessary until timed
            seconds = delta.total_seconds()
            if seconds > 0:
                await asyncio.sleep(seconds)

            # Send reminder
            self.bot.privmsg(target, '\x02[Timer]\x0F {nick}: {message} '
                                     '({delay})'.format(message=message,
                                                        nick=mask.nick,
                                                        delay=delay))

            try:
                # Delete timer from database
                self.cur.execute('''
                delete from
                  timers
                where
                  id = %s
                ''', [row_id])
                self.con.commit()
            except Error:
                # Rollback transaction on error
                self.con.rollback()

        asyncio.ensure_future(callback())