1
Fork 0
mirror of https://github.com/RYGhub/royalnet.git synced 2024-11-27 13:34:28 +00:00
This commit is contained in:
Steffo 2019-11-14 13:50:26 +01:00
parent 756b9d25ed
commit 9176144392
2 changed files with 126 additions and 132 deletions

View file

@ -1,35 +1,50 @@
import discord import logging
import sentry_sdk import asyncio
import logging as _logging from typing import Type, Optional, List, Callable, Union
from .generic import GenericBot from royalnet.commands import Command, CommandInterface, CommandData, CommandArgs, CommandError, InvalidInputError, \
from royalnet.utils import * UnsupportedError
from royalnet.error import * from royalnet.utils import asyncify
from royalnet.bard import * from .escape import escape
from royalnet.commands import * from ..serf import Serf
try:
import discord
except ImportError:
discord = None
try:
from sqlalchemy.orm.session import Session
from ..alchemyconfig import AlchemyConfig
except ImportError:
Session = None
AlchemyConfig = None
try:
from royalnet.herald import Config as HeraldConfig
except ImportError:
HeraldConfig = None
log = logging.getLogger(__name__)
log = _logging.getLogger(__name__) class DiscordSerf(Serf):
"""A :class:`Serf` that connects to `Discord <https://discordapp.com/>`_ as a bot."""
class MusicData:
def __init__(self):
self.playmode: playmodes.PlayMode = playmodes.Playlist()
self.voice_client: typing.Optional[discord.VoiceClient] = None
def queue_preview(self):
return self.playmode.queue_preview()
class DiscordBot(GenericBot):
"""A bot that connects to `Discord <https://discordapp.com/>`_."""
interface_name = "discord" interface_name = "discord"
def _init_voice(self): def __init__(self, *,
"""Initialize the variables needed for the connection to voice chat.""" alchemy_config: Optional[AlchemyConfig] = None,
log.debug(f"Creating music_data dict") commands: List[Type[Command]] = None,
self.music_data: typing.Dict[discord.Guild, MusicData] = {} network_config: Optional[HeraldConfig] = None,
secrets_name: str = "__default__"):
if discord is None:
raise ImportError("'discord' extra is not installed")
def _interface_factory(self) -> typing.Type[CommandInterface]: super().__init__(alchemy_config=alchemy_config,
commands=commands,
network_config=network_config,
secrets_name=secrets_name)
def _interface_factory(self) -> Type[CommandInterface]:
# noinspection PyPep8Naming # noinspection PyPep8Naming
GenericInterface = super().interface_factory() GenericInterface = super().interface_factory()
@ -40,22 +55,22 @@ class DiscordBot(GenericBot):
return DiscordInterface return DiscordInterface
def _data_factory(self) -> typing.Type[CommandData]: def _data_factory(self) -> Type[CommandData]:
# noinspection PyMethodParameters,PyAbstractClass # noinspection PyMethodParameters,PyAbstractClass
class DiscordData(CommandData): class DiscordData(CommandData):
def __init__(data, interface: CommandInterface, message: discord.Message): def __init__(data, interface: CommandInterface, session, message: discord.Message):
super().__init__(interface) super().__init__(interface=interface, session=session)
data.message = message data.message = message
async def reply(data, text: str): async def reply(data, text: str):
await data.message.channel.send(discord_escape(text)) await data.message.channel.send(escape(text))
async def get_author(data, error_if_none=False): async def get_author(data, error_if_none=False):
user: discord.Member = data.message.author user: discord.Member = data.message.author
query = data.session.query(self.master_table) query = data.session.query(self._master_table)
for link in self.identity_chain: for link in self._identity_chain:
query = query.join(link.mapper.class_) query = query.join(link.mapper.class_)
query = query.filter(self.identity_column == user.id) query = query.filter(self._identity_column == user.id)
result = await asyncify(query.one_or_none) result = await asyncify(query.one_or_none)
if result is None and error_if_none: if result is None and error_if_none:
raise CommandError("You must be registered to use this command.") raise CommandError("You must be registered to use this command.")
@ -66,42 +81,8 @@ class DiscordBot(GenericBot):
return DiscordData return DiscordData
def _bot_factory(self) -> typing.Type[discord.Client]: async def handle_message(self, message: discord.Message):
"""Create a custom DiscordClient class inheriting from :py:class:`discord.Client`.""" """Handle a Discord message by calling a command if appropriate."""
log.debug(f"Creating DiscordClient")
# noinspection PyMethodParameters
class DiscordClient(discord.Client):
async def vc_connect_or_move(cli, channel: discord.VoiceChannel):
music_data = self.music_data.get(channel.guild)
if music_data is None:
# Create a MusicData object
music_data = MusicData()
self.music_data[channel.guild] = music_data
# Connect to voice
log.debug(f"Connecting to Voice in {channel}")
try:
music_data.voice_client = await channel.connect(reconnect=False, timeout=10)
except Exception:
log.warning(f"Failed to connect to Voice in {channel}")
del self.music_data[channel.guild]
raise
else:
log.debug(f"Connected to Voice in {channel}")
else:
if music_data.voice_client is None:
# TODO: change exception type
raise Exception("Another connection attempt is already in progress.")
# Try to move to a different channel
voice_client = music_data.voice_client
log.debug(f"Moving {voice_client} to {channel}")
await voice_client.move_to(channel)
log.debug(f"Moved {voice_client} to {channel}")
async def on_message(cli, message: discord.Message):
self.loop.create_task(cli._handle_message(message))
async def _handle_message(cli, message: discord.Message):
text = message.content text = message.content
# Skip non-text messages # Skip non-text messages
if not text: if not text:
@ -110,7 +91,7 @@ class DiscordBot(GenericBot):
if not text.startswith("!"): if not text.startswith("!"):
return return
# Skip bot messages # Skip bot messages
author: typing.Union[discord.User] = message.author author: Union[discord.User] = message.author
if author.bot: if author.bot:
return return
# Find and clean parameters # Find and clean parameters
@ -123,13 +104,18 @@ class DiscordBot(GenericBot):
except KeyError: except KeyError:
# Skip the message # Skip the message
return return
# Prepare data
data = self._Data(interface=command.interface, message=message)
# Call the command # Call the command
log.debug(f"Calling command '{command.name}'") log.debug(f"Calling command '{command.name}'")
with message.channel.typing(): with message.channel.typing():
# Run the command # Open an alchemy session, if available
if self.alchemy is not None:
session = await asyncify(self.alchemy.Session)
else:
session = None
# Prepare data
data = self.Data(interface=command.interface, session=session, message=message)
try: try:
# Run the command
await command.run(CommandArgs(parameters), data) await command.run(CommandArgs(parameters), data)
except InvalidInputError as e: except InvalidInputError as e:
await data.reply(f":warning: {e.message}\n" await data.reply(f":warning: {e.message}\n"
@ -139,44 +125,52 @@ class DiscordBot(GenericBot):
except CommandError as e: except CommandError as e:
await data.reply(f":warning: {e.message}") await data.reply(f":warning: {e.message}")
except Exception as e: except Exception as e:
sentry_sdk.capture_exception(e) self.sentry_exc(e)
error_message = f"🦀 [b]{e.__class__.__name__}[/b] 🦀\n" error_message = f"🦀 [b]{e.__class__.__name__}[/b] 🦀\n" \
error_message += '\n'.join(e.args) '\n'.join(e.args)
await data.reply(error_message) await data.reply(error_message)
# Close the data session finally:
await data.session_close() # Close the alchemy session
if session is not None:
await asyncify(session.close)
async def on_connect(cli): def _bot_factory(self) -> Type[discord.Client]:
log.debug("Connected to Discord") """Create a custom class inheriting from :py:class:`discord.Client`."""
# noinspection PyMethodParameters
async def on_disconnect(cli): class DiscordClient(discord.Client):
log.error("Disconnected from Discord!") async def on_message(cli, message: discord.Message):
"""Handle messages received by passing them to the handle_message method of the bot."""
# TODO: keep reference to these tasks somewhere
self.loop.create_task(self.handle_message(message))
async def on_ready(cli) -> None: async def on_ready(cli) -> None:
log.debug("Connection successful, client is ready") """Change the bot presence to ``online`` when the bot is ready."""
await cli.change_presence(status=discord.Status.online) await cli.change_presence(status=discord.Status.online)
def find_guild_by_name(cli, name: str) -> typing.List[discord.Guild]: def find_guild(cli, name: str) -> List[discord.Guild]:
"""Find the :py:class:`discord.Guild` with the specified name (case insensitive).""" """Find the :class:`discord.Guild`s with the specified name (case insensitive).
all_guilds: typing.List[discord.Guild] = cli.guilds
matching_channels: typing.List[discord.Guild] = [] Returns:
A :class:`list` of :class:`discord.Guild` having the specified name."""
all_guilds: List[discord.Guild] = cli.guilds
matching_channels: List[discord.Guild] = []
for guild in all_guilds: for guild in all_guilds:
if guild.name.lower() == name.lower(): if guild.name.lower() == name.lower():
matching_channels.append(guild) matching_channels.append(guild)
return matching_channels return matching_channels
def find_channel_by_name(cli, def find_channel(cli,
name: str, name: str,
guild: typing.Optional[discord.Guild] = None) -> typing.List[discord.abc.GuildChannel]: guild: Optional[discord.Guild] = None) -> List[discord.abc.GuildChannel]:
"""Find the :py:class:`TextChannel`, :py:class:`VoiceChannel` or :py:class:`CategoryChannel` with the """Find the :class:`TextChannel`, :class:`VoiceChannel` or :class:`CategoryChannel` with the
specified name (case insensitive). specified name (case insensitive).
You can specify a guild to only find channels in that specific guild.""" You can specify a guild to only search in that specific guild."""
if guild is not None: if guild is not None:
all_channels = guild.channels all_channels = guild.channels
else: else:
all_channels: typing.List[discord.abc.GuildChannel] = cli.get_all_channels() all_channels: List[discord.abc.GuildChannel] = cli.get_all_channels()
matching_channels: typing.List[discord.abc.GuildChannel] = [] matching_channels: List[discord.abc.GuildChannel] = []
for channel in all_channels: for channel in all_channels:
if not (isinstance(channel, discord.TextChannel) if not (isinstance(channel, discord.TextChannel)
or isinstance(channel, discord.VoiceChannel) or isinstance(channel, discord.VoiceChannel)
@ -186,8 +180,9 @@ class DiscordBot(GenericBot):
matching_channels.append(channel) matching_channels.append(channel)
return matching_channels return matching_channels
def find_voice_client_by_guild(cli, guild: discord.Guild) -> typing.Optional[discord.VoiceClient]: def find_voice_client(cli, guild: discord.Guild) -> Optional[discord.VoiceClient]:
"""Find the :py:class:`discord.VoiceClient` belonging to a specific :py:class:`discord.Guild`.""" """Find the :py:class:`discord.VoiceClient` belonging to a specific :py:class:`discord.Guild`."""
# TODO: the bug I was looking for might be here
for voice_client in cli.voice_clients: for voice_client in cli.voice_clients:
if voice_client.guild == guild: if voice_client.guild == guild:
return voice_client return voice_client
@ -195,6 +190,8 @@ class DiscordBot(GenericBot):
return DiscordClient return DiscordClient
# TODO: restart from here
def _init_client(self): def _init_client(self):
"""Create an instance of the DiscordClient class created in :py:func:`royalnet.bots.DiscordBot._bot_factory`.""" """Create an instance of the DiscordClient class created in :py:func:`royalnet.bots.DiscordBot._bot_factory`."""
log.debug(f"Creating DiscordClient instance") log.debug(f"Creating DiscordClient instance")

View file

@ -32,7 +32,7 @@ log = logging.getLogger(__name__)
class TelegramSerf(Serf): class TelegramSerf(Serf):
"""A Serf that connects to `Telegram <https://telegram.org/>`_.""" """A Serf that connects to `Telegram <https://telegram.org/>`_ as a bot."""
interface_name = "telegram" interface_name = "telegram"
def __init__(self, *, def __init__(self, *,
@ -94,9 +94,6 @@ class TelegramSerf(Serf):
name = self.interface_name name = self.interface_name
prefix = "/" prefix = "/"
def __init__(self):
super().__init__()
return TelegramInterface return TelegramInterface
def data_factory(self) -> Type[CommandData]: def data_factory(self) -> Type[CommandData]: