python-app-rocketbot/rocketbot/cogs/basecog.py

"""
Base cog class and helper classes.
"""
from datetime import datetime, timedelta
from discord import Guild, Member, Message, RawReactionActionEvent
from discord.abc import GuildChannel
from discord.ext import commands

from config import CONFIG
from rocketbot.botmessage import BotMessage, BotMessageReaction
from rocketbot.cogsetting import CogSetting
from rocketbot.collections import AgeBoundDict
from rocketbot.storage import Storage
from rocketbot.utils import bot_log

class BaseCog(commands.Cog):
	"""
	Superclass for all Rocketbot cogs. Provides lots of conveniences for
	common tasks.
	"""
	def __init__(self, bot):
		self.bot = bot
		self.are_settings_setup = False
		self.settings = []

	# Config

	@classmethod
	def get_cog_default(cls, key: str):
		"""
		Convenience method for getting a cog configuration default from
		`CONFIG['cogs'][<cogname>][<key>]`. These values are used for
		CogSettings when no guild-specific value is configured yet.
		"""
		cogs: dict = CONFIG['cog_defaults']
		cog = cogs.get(cls.__name__)
		if cog is None:
			return None
		return cog.get(key)

	def add_setting(self, setting: CogSetting) -> None:
		"""
		Called by a subclass in __init__ to register a mod-configurable
		guild setting. A "get" and "set" command will be generated. If the
		setting is named "enabled" (exactly) then "enable" and "disable"
		commands will be created instead which set the setting to True/False.

		If the cog has a command group it will be detected automatically and
		the commands added to that. Otherwise the commands will be added at
		the top level.

		Changes to settings can be detected by overriding `on_setting_updated`.
		"""
		self.settings.append(setting)

	@classmethod
	def get_guild_setting(cls,
			guild: Guild,
			setting: CogSetting,
			use_cog_default_if_not_set: bool = True):
		"""
		Returns the configured value for a setting for the given guild. If no
		setting is configured the default for the cog will be returned,
		unless the optional `use_cog_default_if_not_set` is `False`, then
		`None` will be returned.
		"""
		key = f'{cls.__name__}.{setting.name}'
		value = Storage.get_config_value(guild, key)
		if value is None and use_cog_default_if_not_set:
			value = cls.get_cog_default(setting.name)
		return value

	@classmethod
	def set_guild_setting(cls,
			guild: Guild,
			setting: CogSetting,
			new_value) -> None:
		"""
		Manually sets a setting for the given guild. BaseCog creates "get" and
		"set" commands for guild administrators to configure values themselves,
		but this method can be used for hidden settings from code. A ValueError
		will be raised if the new value does not pass validation specified in
		the CogSetting.
		"""
		setting.validate_value(new_value)
		key = f'{cls.__name__}.{setting.name}'
		Storage.set_config_value(guild, key, new_value)

	@commands.Cog.listener()
	async def on_ready(self):
		'Event listener'
		if not self.are_settings_setup:
			self.are_settings_setup = True
			CogSetting.set_up_all(self, self.bot, self.settings)

	async def on_setting_updated(self, guild: Guild, setting: CogSetting) -> None:
		"""
		Subclass override point for being notified when a CogSetting is edited.
		"""

	# Bot message handling

	@classmethod
	def __bot_messages(cls, guild: Guild) -> AgeBoundDict:
		bm = Storage.get_state_value(guild, 'bot_messages')
		if bm is None:
			far_future = datetime.utcnow() + timedelta(days=1000)
			bm = AgeBoundDict(timedelta(seconds=600),
				lambda k, v : v.message_sent_at() or far_future)
			Storage.set_state_value(guild, 'bot_messages', bm)
		return bm

	async def post_message(self, message: BotMessage) -> bool:
		"""
		Posts a BotMessage to a guild. Returns whether it was successful. If
		the caller wants to listen to reactions they should be added before
		calling this method. Listen to reactions by overriding `on_mod_react`.
		"""
		message.source_cog = self
		await message.update()
		if message.has_reactions() and message.is_sent():
			guild_messages = self.__bot_messages(message.guild)
			guild_messages[message.message_id()] = message
		return message.is_sent()

	@commands.Cog.listener()
	async def on_raw_reaction_add(self, payload: RawReactionActionEvent):
		'Event handler'
		if payload.user_id == self.bot.user.id:
			# Ignore bot's own reactions
			return
		member: Member = payload.member
		if member is None:
			return
		guild: Guild = self.bot.get_guild(payload.guild_id)
		if guild is None:
			# Possibly a DM
			return
		channel: GuildChannel = guild.get_channel(payload.channel_id)
		if channel is None:
			# Possibly a DM
			return
		message: Message = await channel.fetch_message(payload.message_id)
		if message is None:
			# Message deleted?
			return
		if message.author.id != self.bot.user.id:
			# Bot didn't author this
			return
		guild_messages = self.__bot_messages(guild)
		bot_message = guild_messages.get(message.id)
		if bot_message is None:
			# Unknown message (expired or was never tracked)
			return
		if self is not bot_message.source_cog:
			# Belongs to a different cog
			return
		reaction = bot_message.reaction_for_emoji(payload.emoji)
		if reaction is None or not reaction.is_enabled:
			# Can't use this reaction with this message
			return
		if not member.permissions_in(channel).ban_members:
			# Not a mod (could make permissions configurable per BotMessageReaction some day)
			return
		await self.on_mod_react(bot_message, reaction, member)

	async def on_mod_react(self,
			bot_message: BotMessage,
			reaction: BotMessageReaction,
			reacted_by: Member) -> None:
		"""
		Subclass override point for receiving mod reactions to bot messages sent
		via `post_message()`.
		"""

	# Helpers

	@classmethod
	def log(cls, guild: Guild, message) -> None:
		"""
		Writes a message to the console. Intended for significant events only.
		"""
		bot_log(guild, cls, message)