pgbot/pgbotlib/cron.py

import asyncio
import random
import sys

import aiocron
import telethon
import pgbotlib.defaults
import pgbotlib.response
import pgbotlib.misc
import pytz
import yaml


class Cron:
    def __init__(self,
                 config: dict,
                 client: telethon.TelegramClient,
                 responder: pgbotlib.response.Responder) -> None:
        schedule_conf = config.get('schedule', pgbotlib.defaults.SCHEDULE)
        local_tz = config.get('timezone', pgbotlib.defaults.TZ)
        try:
            self.tz = pytz.timezone(local_tz)
        except pytz.exceptions.UnknownTimeZoneError as e:
            sys.stderr.write(e)
            self.tz = pytz.utc
        with open(schedule_conf, 'r', encoding='utf-8') as data:
            self.sched = yaml.safe_load(data.read())
            self.responder = responder
            self.client = client

    def __mkjob(self, job: dict) -> callable:
        tokens = frozenset(job['tokens'].split(','))
        async def send_message() -> None:
            if 'rand' in job:
                wait_seconds = random.randint(0, job['rand']) * 60
                await asyncio.sleep(wait_seconds)
            message = self.responder.get_response(tokens)
            message = self.responder.api_match(message, '')
            await self.client.send_message(job['chat'], message)
        return send_message

    def plan(self) -> None:
        for job in self.sched:
            aiocron.crontab(job['cron'], tz=self.tz, func=self.__mkjob(job))