nxy/bot/timer.py

97 lines
2.9 KiB
Python
Raw Normal View History

2017-05-16 10:06:29 +00:00
# -*- coding: utf-8 -*-
import re
2017-05-16 10:06:29 +00:00
import asyncio
from datetime import datetime
2017-05-16 10:06:29 +00:00
from aiocron import crontab
2017-08-22 13:57:57 +00:00
from docopt import Dict
2017-08-22 14:00:03 +00:00
from irc3.plugins.command import command
2017-05-16 10:06:29 +00:00
from irc3.utils import IrcString
from psycopg2 import Error
from psycopg2.extras import DictRow
2017-05-16 10:06:29 +00:00
2017-08-22 15:43:48 +00:00
from . import DatabasePlugin, Bot
2017-05-16 10:06:29 +00:00
class Timer(DatabasePlugin):
2017-08-22 15:43:48 +00:00
def __init__(self, bot: Bot):
2017-05-16 10:06:29 +00:00
super().__init__(bot)
self.timers = set()
self.set_timers()
crontab('0 * * * *', func=self.set_timers)
2017-06-30 12:04:50 +00:00
2017-05-16 10:06:29 +00:00
@command
2017-08-22 13:57:57 +00:00
def timer(self, mask: IrcString, target: IrcString, args: Dict):
2017-08-16 12:39:44 +00:00
"""Sets a timer, delay can be: s, m, h, d, w, mon, y
2017-05-16 11:44:39 +00:00
2017-05-16 10:06:29 +00:00
%%timer <delay> <message>...
"""
delay = args['<delay>']
message = ' '.join(args['<message>'])
2017-08-22 11:33:50 +00:00
# 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()
2017-08-22 11:22:57 +00:00
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))
2017-08-22 11:22:57 +00:00
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()