-
-
Notifications
You must be signed in to change notification settings - Fork 9
Bump d.py and add bot base #42
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
12 commits
Select commit
Hold shift + click to select a range
a6ab7b1
Bump d.py version and add new requirements for utils
ChrisLovering 1e6556a
Add async_stats subclass of statsd
ChrisLovering 8bd3706
Add BotBase that will act as a base for all our bots
ChrisLovering 9079605
Remove unneeded kwargs in BotBase
ChrisLovering b8a2450
Provide a bot.all_extensions instance attribute
ChrisLovering 08d4b71
Comment on what nitpick_ignore_regex does in conf.py
ChrisLovering 2383754
Mark async-rediscache as an optional extra dependency
ChrisLovering 3e71424
Remove inappropriate abstractmethod tags
ChrisLovering 8e1e881
Add __all__ attrs to async_stats and site_api modules
ChrisLovering 5a02b3b
Remove BotBase.closing_tasks
ChrisLovering 38f0e15
Install optional deps during CI
ChrisLovering 7d00aec
Use imperative mood in docstrings
ChrisLovering File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,11 +1,15 @@ | ||
"""Useful utilities and tools for Discord bot development.""" | ||
|
||
from botcore import exts, site_api, utils | ||
from botcore import async_stats, exts, site_api, utils | ||
from botcore._bot import BotBase, StartupError | ||
|
||
__all__ = [ | ||
async_stats, | ||
BotBase, | ||
exts, | ||
utils, | ||
site_api, | ||
MarkKoz marked this conversation as resolved.
Show resolved
Hide resolved
|
||
StartupError, | ||
] | ||
|
||
__all__ = list(map(lambda module: module.__name__, __all__)) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,262 @@ | ||
import asyncio | ||
import socket | ||
import types | ||
import warnings | ||
from contextlib import suppress | ||
from typing import Optional | ||
|
||
import aiohttp | ||
import discord | ||
from discord.ext import commands | ||
|
||
from botcore.async_stats import AsyncStatsClient | ||
from botcore.site_api import APIClient | ||
from botcore.utils._extensions import walk_extensions | ||
from botcore.utils.logging import get_logger | ||
|
||
try: | ||
from async_rediscache import RedisSession | ||
except ImportError: | ||
RedisSession = discord.utils._MissingSentinel | ||
|
||
log = get_logger() | ||
|
||
|
||
class StartupError(Exception): | ||
"""Exception class for startup errors.""" | ||
|
||
def __init__(self, base: Exception): | ||
super().__init__() | ||
self.exception = base | ||
|
||
|
||
class BotBase(commands.Bot): | ||
"""A sub-class that implements many common features that Python Discord bots use.""" | ||
|
||
def __init__( | ||
self, | ||
*args, | ||
guild_id: int, | ||
allowed_roles: list, | ||
http_session: aiohttp.ClientSession, | ||
redis_session: Optional[RedisSession] = None, | ||
**kwargs, | ||
): | ||
""" | ||
Initialise the base bot instance. | ||
|
||
Args: | ||
guild_id: The ID of the guild use for :func:`wait_until_guild_available`. | ||
allowed_roles: A list of role IDs that the bot is allowed to mention. | ||
http_session (aiohttp.ClientSession): The session to use for the bot. | ||
redis_session: The | ||
``[async_rediscache.RedisSession](https://github.com/SebastiaanZ/async-rediscache#creating-a-redissession)`` | ||
to use for the bot. | ||
""" | ||
super().__init__( | ||
*args, | ||
allowed_roles=allowed_roles, | ||
**kwargs, | ||
) | ||
|
||
self.guild_id = guild_id | ||
self.http_session = http_session | ||
|
||
if redis_session and RedisSession == discord.utils._MissingSentinel: | ||
warnings.warn("redis_session kwarg passed, but async-rediscache not installed!") | ||
elif redis_session: | ||
self.redis_session = redis_session | ||
|
||
self.api_client: Optional[APIClient] = None | ||
|
||
self._resolver = aiohttp.AsyncResolver() | ||
self._connector = aiohttp.TCPConnector( | ||
resolver=self._resolver, | ||
family=socket.AF_INET, | ||
) | ||
self.http.connector = self._connector | ||
|
||
self.statsd_url: Optional[str] = None | ||
self._statsd_timerhandle: Optional[asyncio.TimerHandle] = None | ||
self._guild_available = asyncio.Event() | ||
|
||
self.stats: Optional[AsyncStatsClient] = None | ||
|
||
self.all_extensions: Optional[frozenset[str]] = None | ||
|
||
def _connect_statsd( | ||
self, | ||
statsd_url: str, | ||
loop: asyncio.AbstractEventLoop, | ||
retry_after: int = 2, | ||
attempt: int = 1 | ||
) -> None: | ||
"""Callback used to retry a connection to statsd if it should fail.""" | ||
if attempt >= 8: | ||
log.error("Reached 8 attempts trying to reconnect AsyncStatsClient. Aborting") | ||
return | ||
|
||
try: | ||
self.stats = AsyncStatsClient(loop, statsd_url, 8125, prefix="bot") | ||
except socket.gaierror: | ||
log.warning(f"Statsd client failed to connect (Attempt(s): {attempt})") | ||
# Use a fallback strategy for retrying, up to 8 times. | ||
self._statsd_timerhandle = loop.call_later( | ||
retry_after, | ||
self._connect_statsd, | ||
statsd_url, | ||
retry_after * 2, | ||
attempt + 1 | ||
) | ||
|
||
async def load_extensions(self, module: types.ModuleType) -> None: | ||
"""Load all the extensions within the given module and save them to ``self.all_extensions``.""" | ||
self.all_extensions = walk_extensions(module) | ||
ChrisLovering marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
for extension in self.all_extensions: | ||
await self.load_extension(extension) | ||
|
||
def _add_root_aliases(self, command: commands.Command) -> None: | ||
"""Recursively add root aliases for ``command`` and any of its subcommands.""" | ||
if isinstance(command, commands.Group): | ||
for subcommand in command.commands: | ||
self._add_root_aliases(subcommand) | ||
|
||
for alias in getattr(command, "root_aliases", ()): | ||
if alias in self.all_commands: | ||
raise commands.CommandRegistrationError(alias, alias_conflict=True) | ||
|
||
self.all_commands[alias] = command | ||
|
||
def _remove_root_aliases(self, command: commands.Command) -> None: | ||
"""Recursively remove root aliases for ``command`` and any of its subcommands.""" | ||
if isinstance(command, commands.Group): | ||
for subcommand in command.commands: | ||
self._remove_root_aliases(subcommand) | ||
|
||
for alias in getattr(command, "root_aliases", ()): | ||
self.all_commands.pop(alias, None) | ||
|
||
async def add_cog(self, cog: commands.Cog) -> None: | ||
"""Add the given ``cog`` to the bot and log the operation.""" | ||
await super().add_cog(cog) | ||
log.info(f"Cog loaded: {cog.qualified_name}") | ||
|
||
def add_command(self, command: commands.Command) -> None: | ||
"""Add ``command`` as normal and then add its root aliases to the bot.""" | ||
super().add_command(command) | ||
self._add_root_aliases(command) | ||
|
||
def remove_command(self, name: str) -> Optional[commands.Command]: | ||
""" | ||
Remove a command/alias as normal and then remove its root aliases from the bot. | ||
|
||
Individual root aliases cannot be removed by this function. | ||
To remove them, either remove the entire command or manually edit `bot.all_commands`. | ||
""" | ||
command = super().remove_command(name) | ||
if command is None: | ||
# Even if it's a root alias, there's no way to get the Bot instance to remove the alias. | ||
return None | ||
|
||
self._remove_root_aliases(command) | ||
return command | ||
|
||
def clear(self) -> None: | ||
"""Not implemented! Re-instantiate the bot instead of attempting to re-use a closed one.""" | ||
raise NotImplementedError("Re-using a Bot object after closing it is not supported.") | ||
|
||
async def on_guild_unavailable(self, guild: discord.Guild) -> None: | ||
"""Clear the internal guild available event when self.guild_id becomes unavailable.""" | ||
if guild.id != self.guild_id: | ||
return | ||
|
||
self._guild_available.clear() | ||
|
||
async def on_guild_available(self, guild: discord.Guild) -> None: | ||
""" | ||
Set the internal guild available event when self.guild_id becomes available. | ||
|
||
If the cache appears to still be empty (no members, no channels, or no roles), the event | ||
will not be set and `guild_available_but_cache_empty` event will be emitted. | ||
""" | ||
if guild.id != self.guild_id: | ||
return | ||
|
||
if not guild.roles or not guild.members or not guild.channels: | ||
msg = "Guild available event was dispatched but the cache appears to still be empty!" | ||
self.log_to_dev_log(msg) | ||
return | ||
|
||
self._guild_available.set() | ||
|
||
async def log_to_dev_log(self, message: str) -> None: | ||
"""Log the given message to #dev-log.""" | ||
... | ||
|
||
async def wait_until_guild_available(self) -> None: | ||
""" | ||
Wait until the guild that matches the ``guild_id`` given at init is available (and the cache is ready). | ||
|
||
The on_ready event is inadequate because it only waits 2 seconds for a GUILD_CREATE | ||
gateway event before giving up and thus not populating the cache for unavailable guilds. | ||
""" | ||
await self._guild_available.wait() | ||
|
||
async def setup_hook(self) -> None: | ||
""" | ||
An async init to startup generic services. | ||
|
||
Connects to statsd, and calls | ||
:func:`AsyncStatsClient.create_socket <botcore.async_stats.AsyncStatsClient.create_socket>` | ||
and :func:`ping_services`. | ||
""" | ||
loop = asyncio.get_running_loop() | ||
|
||
self._connect_statsd(self.statsd_url, loop) | ||
self.stats = AsyncStatsClient(loop, "127.0.0.1") | ||
await self.stats.create_socket() | ||
|
||
try: | ||
await self.ping_services() | ||
except Exception as e: | ||
raise StartupError(e) | ||
|
||
async def ping_services() -> None: | ||
"""Ping all required services on setup to ensure they are up before starting.""" | ||
... | ||
|
||
async def close(self) -> None: | ||
"""Close the Discord connection, and the aiohttp session, connector, statsd client, and resolver.""" | ||
# Done before super().close() to allow tasks finish before the HTTP session closes. | ||
for ext in list(self.extensions): | ||
MarkKoz marked this conversation as resolved.
Show resolved
Hide resolved
|
||
with suppress(Exception): | ||
await self.unload_extension(ext) | ||
|
||
for cog in list(self.cogs): | ||
with suppress(Exception): | ||
await self.remove_cog(cog) | ||
|
||
# Now actually do full close of bot | ||
await super().close() | ||
|
||
if self.api_client: | ||
await self.api_client.close() | ||
|
||
if self.http_session: | ||
await self.http_session.close() | ||
|
||
if self._connector: | ||
await self._connector.close() | ||
|
||
if self._resolver: | ||
await self._resolver.close() | ||
|
||
if self.stats._transport: | ||
self.stats._transport.close() | ||
|
||
if getattr(self.redis_session, None): | ||
await self.redis_session.close() | ||
|
||
if self._statsd_timerhandle: | ||
self._statsd_timerhandle.cancel() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,57 @@ | ||
"""An async transport method for statsd communication.""" | ||
|
||
import asyncio | ||
import socket | ||
from typing import Optional | ||
|
||
from statsd.client.base import StatsClientBase | ||
|
||
from botcore.utils import scheduling | ||
|
||
|
||
class AsyncStatsClient(StatsClientBase): | ||
"""An async implementation of :obj:`statsd.client.base.StatsClientBase` that supports async stat communication.""" | ||
|
||
def __init__( | ||
self, | ||
loop: asyncio.AbstractEventLoop, | ||
host: str = 'localhost', | ||
port: int = 8125, | ||
prefix: str = None | ||
): | ||
""" | ||
Create a new :obj:`AsyncStatsClient`. | ||
|
||
Args: | ||
loop (asyncio.AbstractEventLoop): The event loop to use when creating the | ||
:obj:`asyncio.loop.create_datagram_endpoint`. | ||
host: The host to connect to. | ||
port: The port to connect to. | ||
prefix: The prefix to use for all stats. | ||
""" | ||
_, _, _, _, addr = socket.getaddrinfo( | ||
host, port, socket.AF_INET, socket.SOCK_DGRAM | ||
)[0] | ||
self._addr = addr | ||
self._prefix = prefix | ||
self._loop = loop | ||
self._transport: Optional[asyncio.DatagramTransport] = None | ||
|
||
async def create_socket(self) -> None: | ||
"""Use :obj:`asyncio.loop.create_datagram_endpoint` from the loop given on init to create a socket.""" | ||
self._transport, _ = await self._loop.create_datagram_endpoint( | ||
asyncio.DatagramProtocol, | ||
family=socket.AF_INET, | ||
remote_addr=self._addr | ||
) | ||
|
||
def _send(self, data: str) -> None: | ||
"""Start an async task to send data to statsd.""" | ||
scheduling.create_task(self._async_send(data), event_loop=self._loop) | ||
|
||
async def _async_send(self, data: str) -> None: | ||
"""Send data to the statsd server using the async transport.""" | ||
self._transport.sendto(data.encode('ascii'), self._addr) | ||
|
||
|
||
__all__ = ['AsyncStatsClient'] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Binary file not shown.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.