86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
import discord
|
|
from asyncio import Queue, create_task
|
|
import concurrent.futures
|
|
|
|
"""
|
|
Abstraction around the Discord client library
|
|
"""
|
|
|
|
|
|
class DiscordClient(discord.Client):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._commands = {}
|
|
self._queue = Queue()
|
|
self.bg_task = self.loop.create_task(self._send_loop())
|
|
|
|
def add_command(self, command, responder):
|
|
"""
|
|
Add a command.
|
|
|
|
:param command: The command prefix. For example, '!command'.
|
|
:param responder: The function that responds to the message. This can return either a string, an embed or None.
|
|
"""
|
|
self._commands[command] = responder
|
|
|
|
async def on_message(self, message):
|
|
"""
|
|
Called whenever a message is sent to a Discord channel that ButlerBot can read.
|
|
This overrides on_message() from discord.py.
|
|
|
|
:param message: The message sent by the user.
|
|
"""
|
|
# Don't listen to messages coming from ButlerBot.
|
|
if message.author == self.user:
|
|
return
|
|
|
|
# Check if message matches any of the supported command prefixes.
|
|
for command, responder in self._commands.items():
|
|
command_prefix = message.content.split(' ')[0]
|
|
if command_prefix == command:
|
|
await self._do_command(responder, message)
|
|
|
|
async def _do_command(self, responder, message):
|
|
"""
|
|
Respond to a command.
|
|
|
|
:param responder: The function that responds to the message.
|
|
:param message: The message sent by the user.
|
|
"""
|
|
# Some commands take a while. Let users know by showing the Discord typing indicator.
|
|
with concurrent.futures.ThreadPoolExecutor() as pool:
|
|
response_placeholder = await message.channel.send(content=":hourglass: **Command in progress...**")
|
|
await message.channel.trigger_typing()
|
|
response = await self.loop.run_in_executor(pool, responder, message)
|
|
await response_placeholder.delete()
|
|
if response is None:
|
|
pass
|
|
elif isinstance(response, discord.Embed):
|
|
await message.channel.send(embed=response)
|
|
else:
|
|
await message.channel.send(content=response)
|
|
|
|
def queue_message(self, channel, message):
|
|
"""
|
|
Queues a message to be sent.
|
|
|
|
:param channel: The channel in which the message will be sent.
|
|
:param message: The message to be sent. This can be eiher a string or an embed.
|
|
"""
|
|
self._queue.put_nowait((channel, message))
|
|
|
|
async def _send_loop(self):
|
|
"""
|
|
Sends queued messages.
|
|
"""
|
|
await self.wait_until_ready()
|
|
while True:
|
|
channel, message = await self._queue.get()
|
|
if channel is not None and message is not None:
|
|
if isinstance(message, discord.Embed):
|
|
await channel.send(embed=message)
|
|
else:
|
|
await channel.send(content=message)
|
|
self._queue.task_done()
|