Edm0nd/edmond/plugins/sleep.py

import random
from datetime import datetime, timedelta

from edmond.plugin import Plugin
from edmond.utils import proc


class SleepPlugin(Plugin):
    """Handle sleep state of the bot, snore a little bit."""

    REQUIRED_CONFIGS = [
        "commands", "sleep_time", "wakeup_time", "snore", "sleep_messages",
        "wakeup_messages", "snore_rate"
    ]

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

    def on_welcome(self, event):
        self.set_runtime_value("awake", True)

    def on_ping(self, event):
        awake = self.get_runtime_value("awake")
        in_sleep_hours = self.is_sleep_time(datetime.now())
        if awake and in_sleep_hours:
            self.fall_asleep()
        elif not awake:
            if not in_sleep_hours:
                self.wake_up()
            else:
                # Maybe snore.
                if proc(self.config["snore_rate"]):
                    for channel in self.bot.channels:
                        self.bot.say(channel, self.config["snore"])

    def on_pubmsg(self, event):
        if not self.should_handle_command(
            event.arguments[0],
            no_content=True,
            exclude_conditions=["sleep"],
        ):
            return False

        # "sleep" command.
        if self.command.ident == self.config["commands"][0]:
            self.fall_asleep()
            return True

        # "wake up" command.
        if self.command.ident == self.config["commands"][1]:
            self.wake_up()
            return True

        return False

    def fall_asleep(self):
        if not self.get_runtime_value("awake"):
            return
        for channel in self.bot.channels:
            self.bot.say(channel, random.choice(self.config["sleep_messages"]))
        self.set_runtime_value("awake", False)

    def wake_up(self):
        if self.get_runtime_value("awake"):
            return
        self.set_runtime_value("awake", True)
        for channel in self.bot.channels:
            self.bot.say(channel, random.choice(self.config["wakeup_messages"]))

    def is_sleep_time(self, now):
        """Return True if the bot should be sleeping by now.

        The sleep range can span over 2 days (e.g. 23 to 7) or be contained in a
        day (e.g. 0 to 8).
        """
        sleep_time = self.config["sleep_time"]
        wakeup_time = self.config["wakeup_time"]
        current_hour = now.hour
        # If we're before the sleep hour, check that we are not in the
        # previous sleep range. Else just check today's range.
        if current_hour < sleep_time:
            ref_dt = now - timedelta(days=1)
        else:
            ref_dt = now
        sleep_dt, wakeup_dt = self._get_sleep_range(
            ref_dt,
            sleep_time,
            wakeup_time
        )
        return sleep_dt <= now < wakeup_dt

    @staticmethod
    def _get_sleep_range(now, sleep_time, wakeup_time):
        """Return the (start, end) datetime tuple using config values.

        Properly accounts for a sleep range going around midnight.
        """
        sleep_dt = now.replace(hour=sleep_time, minute=0, second=0)
        wakeup_dt = now.replace(hour=wakeup_time, minute=0, second=0)
        if wakeup_time < sleep_time:
            wakeup_dt += timedelta(days=1)
        return sleep_dt, wakeup_dt