From 9d8405092eecadf85988cd17eb7dcafee26d3c25 Mon Sep 17 00:00:00 2001 From: mrhanky Date: Tue, 22 Aug 2017 13:20:44 +0200 Subject: [PATCH] Fixed timer plugin, now checks every hour (with aiocron) for timers in the next hour and sets them --- bot/plugins/timer.py | 122 ++++++++++++++++++++----------------------- bot/utils.py | 45 ++-------------- requirements.txt | 1 + 3 files changed, 63 insertions(+), 105 deletions(-) diff --git a/bot/plugins/timer.py b/bot/plugins/timer.py index 089bcb4..e1cabb3 100644 --- a/bot/plugins/timer.py +++ b/bot/plugins/timer.py @@ -1,15 +1,17 @@ # -*- coding: utf-8 -*- +import re import asyncio -from datetime import datetime, timedelta +from datetime import datetime import irc3 +from aiocron import crontab from docopt import Dict as DocOptDict from irc3.plugins.command import command from irc3.utils import IrcString from psycopg2 import Error +from psycopg2.extras import DictRow from . import DatabasePlugin -from ..utils import time_delta @irc3.plugin @@ -19,23 +21,25 @@ class Timer(DatabasePlugin): def __init__(self, bot: irc3.IrcBot): super().__init__(bot) + self.timers = set() + self.set_timers() + crontab('0 * * * *', func=self.set_timers) - # Fetch timers from database + 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 - id, mask, target, message, delay, ends_at - from + SELECT + * + FROM timers + WHERE + ends_at >= now() + AND ends_at < now() + INTERVAL '1h' ''') - # 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']) + for timer in self.cur.fetchall(): + asyncio.ensure_future(self.exec_timer(timer)) @command def timer(self, mask: IrcString, target: IrcString, args: DocOptDict): @@ -44,63 +48,51 @@ class Timer(DatabasePlugin): %%timer ... """ delay = args[''] - delta = time_delta(delay) + message = ' '.join(args['']) - if not delta: - self.bot.privmsg(target, 'Invalid timer delay') - else: - message = ' '.join(args['']) - values = [mask, target, message, delay] + if not re.match(r'\d+[smhdwy]|mon', delay): + return 'Invalid timer delay: {}'.format(delay) - try: - # Insert into database (add now + delta) - self.cur.execute(''' - insert into + values = [mask, target, message, delay] + + try: + 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() + VALUES + (%s, %s, %s, %s, now() + INTERVAL %s) + RETURNING * + ''', values + [delay]) + self.con.commit() - # Add delta and id from inserted and start timer - values.extend([delta, self.cur.fetchone()['id']]) - self.start_timer(*values) + asyncio.ensure_future(self.exec_timer(self.cur.fetchone())) - # 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() + 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 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 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 - async def callback(): - # Sleep if necessary until timed - seconds = delta.total_seconds() - if seconds > 0: - await asyncio.sleep(seconds) + self.timers.add(timer['id']) + seconds = (timer['ends_at'] - datetime.now()).total_seconds() + if seconds > 0.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)) + self.bot.privmsg(timer['target'], '\x02[Timer]\x02 {nick}: {message} ({delay})'.format( + message=timer['message'], + nick=IrcString(timer['mask']).nick, + delay=timer['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()) + self.timers.remove(timer['id']) + self.cur.execute(''' + DELETE FROM + timers + WHERE + id = %s + ''', [timer['id']]) + self.con.commit() diff --git a/bot/utils.py b/bot/utils.py index a34fef5..5142aa3 100644 --- a/bot/utils.py +++ b/bot/utils.py @@ -1,51 +1,20 @@ # -*- coding: utf-8 -*- -import re import random from typing import Tuple from pprint import pprint -from datetime import datetime, timedelta - -TIME_UNITS = { - 's': 'seconds', - 'm': 'minutes', - 'h': 'hours', - 'd': 'days', - 'w': 'weeks', - 'mon': 'months', - 'y': 'years', -} +from datetime import datetime pp = pprint +def re_generator(low: int = 5, high: int = 20) -> str: + return 'R{}'.format('E' * random.randint(low, high)) + + def date_from_iso(date: str) -> datetime: return datetime.strptime(date, '%Y-%m-%dT%H:%M:%S.%fZ') -def time_delta(text: str) -> timedelta: - match = re.match(r'(\d+)(s|m|h|d|mon|w|y)', text) - if match: - num, unit = match.groups() - num = int(num) - if unit == 's': - unit = 'seconds' - elif unit == 'm': - unit = 'minutes' - elif unit == 'h': - unit = 'hours' - elif unit == 'd': - unit = 'days' - elif unit == 'w': - unit = 'weeks' - elif unit == 'mon': - unit = 'weeks' - num *= 4 - elif unit == 'y': - unit = 'weeks' - num *= 52 - return timedelta(**{unit: num}) - - def parse_int(val: str, select: bool = True) -> Tuple[int, str]: try: val = int(val) @@ -60,7 +29,3 @@ def parse_int(val: str, select: bool = True) -> Tuple[int, str]: return val, order except ValueError: pass - - -def re_generator(low: int = 5, high: int = 20) -> str: - return 'R{}'.format('E' * random.randint(low, high)) diff --git a/requirements.txt b/requirements.txt index cdd86d6..d486598 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ #git+https://github.com/gawel/irc3.git#egg=irc3 git+https://github.com/mrhanky17/irc3.git#egg=irc3 git+https://github.com/mrhanky17/docopt.git#egg=docopt +aiocron==0.6 psycopg2==2.7.1 requests==2.14.2 feedparser==5.2.1