Skip to content

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 12 commits into from
Apr 2, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/docs.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,11 @@ jobs:
- uses: actions/checkout@v2

- name: Install Python Dependencies
uses: HassanAbouelela/actions/setup-python@setup-python_v1.0.0
uses: HassanAbouelela/actions/setup-python@setup-python_v1.1.0
with:
dev: true
python_version: 3.9
install_args: "--extras async-rediscache"

# Undeclared dependency for `releases`... whoops
# https://github.com/bitprophet/releases/pull/82
Expand Down
3 changes: 2 additions & 1 deletion .github/workflows/lint.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,12 @@ jobs:

steps:
- name: Install Python Dependencies
uses: HassanAbouelela/actions/setup-python@setup-python_v1.0.0
uses: HassanAbouelela/actions/setup-python@setup-python_v1.1.0
with:
# Set dev=true to run pre-commit which is a dev dependency
dev: true
python_version: 3.9
install_args: "--extras async-rediscache"

# We will not run `flake8` here, as we will use a separate flake8
# action. As pre-commit does not support user installs, we set
Expand Down
6 changes: 5 additions & 1 deletion botcore/__init__.py
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,
StartupError,
]

__all__ = list(map(lambda module: module.__name__, __all__))
262 changes: 262 additions & 0 deletions botcore/_bot.py
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)

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):
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()
57 changes: 57 additions & 0 deletions botcore/async_stats.py
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']
3 changes: 3 additions & 0 deletions botcore/site_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,3 +151,6 @@ async def delete(self, endpoint: str, *, raise_for_status: bool = True, **kwargs

await self.maybe_raise_for_status(resp, raise_for_status)
return await resp.json()


__all__ = ['APIClient', 'ResponseCodeError']
3 changes: 1 addition & 2 deletions botcore/utils/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""Useful utilities and tools for Discord bot development."""

from botcore.utils import _monkey_patches, caching, channel, extensions, logging, members, regex, scheduling
from botcore.utils import _monkey_patches, caching, channel, logging, members, regex, scheduling


def apply_monkey_patches() -> None:
Expand All @@ -23,7 +23,6 @@ def apply_monkey_patches() -> None:
apply_monkey_patches,
caching,
channel,
extensions,
logging,
members,
regex,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ def unqualify(name: str) -> str:

def walk_extensions(module: types.ModuleType) -> frozenset[str]:
"""
Yield extension names from the given module.
Return all extension names from the given module.

Args:
module (types.ModuleType): The module to look for extensions in.
Expand Down
Binary file added docs/_static/statsd_additional_objects.inv
Binary file not shown.
Loading