Allow pre_invoke to be used by 3rd party cogs safely. (#3369)

* Okay, so there's a lot in this diff

* fix docstrings

* meh

* fix misleading var name

* meh...

* useful typehints

* Apply suggestions from code review

Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com>

* dep warn in locations suitable

* Fix this...

* Apply suggestions from code review

Co-Authored-By: jack1142 <6032823+jack1142@users.noreply.github.com>

Co-authored-by: jack1142 <6032823+jack1142@users.noreply.github.com>
This commit is contained in:
Michael H 2020-01-15 20:44:21 -05:00 committed by GitHub
parent 27e6f677e8
commit 60dc54b081
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 74 additions and 9 deletions

View File

@ -10,7 +10,19 @@ from datetime import datetime
from enum import IntEnum from enum import IntEnum
from importlib.machinery import ModuleSpec from importlib.machinery import ModuleSpec
from pathlib import Path from pathlib import Path
from typing import Optional, Union, List, Dict, NoReturn from typing import (
Optional,
Union,
List,
Dict,
NoReturn,
Set,
Coroutine,
TypeVar,
Callable,
Awaitable,
Any,
)
from types import MappingProxyType from types import MappingProxyType
import discord import discord
@ -36,6 +48,9 @@ __all__ = ["RedBase", "Red", "ExitCodes"]
NotMessage = namedtuple("NotMessage", "guild") NotMessage = namedtuple("NotMessage", "guild")
PreInvokeCoroutine = Callable[[commands.Context], Awaitable[Any]]
T_BIC = TypeVar("T_BIC", bound=PreInvokeCoroutine)
def _is_submodule(parent, child): def _is_submodule(parent, child):
return parent == child or child.startswith(parent + ".") return parent == child or child.startswith(parent + ".")
@ -150,6 +165,64 @@ class RedBase(commands.GroupMixin, commands.bot.BotBase, RPCMixin): # pylint: d
self._permissions_hooks: List[commands.CheckPredicate] = [] self._permissions_hooks: List[commands.CheckPredicate] = []
self._red_ready = asyncio.Event() self._red_ready = asyncio.Event()
self._red_before_invoke_objs: Set[PreInvokeCoroutine] = set()
@property
def _before_invoke(self): # DEP-WARN
return self._red_before_invoke_method
@_before_invoke.setter
def _before_invoke(self, val): # DEP-WARN
"""Prevent this from being overwritten in super().__init__"""
pass
async def _red_before_invoke_method(self, ctx):
await self.wait_until_red_ready()
return_exceptions = isinstance(ctx.command, commands.commands._AlwaysAvailableCommand)
if self._red_before_invoke_objs:
await asyncio.gather(
*(coro(ctx) for coro in self._red_before_invoke_objs),
return_exceptions=return_exceptions,
)
def remove_before_invoke_hook(self, coro: PreInvokeCoroutine) -> None:
"""
Functional method to remove a `before_invoke` hook.
"""
self._red_before_invoke_objs.discard(coro)
def before_invoke(self, coro: T_BIC) -> T_BIC:
"""
Overridden decorator method for Red's ``before_invoke`` behavior.
This can safely be used purely functionally as well.
3rd party cogs should remove any hooks which they register at unload
using `remove_before_invoke_hook`
Below behavior shared with discord.py:
.. note::
The ``before_invoke`` hooks are
only called if all checks and argument parsing procedures pass
without error. If any check or argument parsing procedures fail
then the hooks are not called.
Parameters
----------
coro: Callable[[commands.Context], Awaitable[Any]]
The coroutine to register as the pre-invoke hook.
Raises
------
TypeError
The coroutine passed is not actually a coroutine.
"""
if not asyncio.iscoroutinefunction(coro):
raise TypeError("The pre-invoke hook must be a coroutine.")
self._red_before_invoke_objs.add(coro)
return coro
@property @property
def cog_mgr(self) -> NoReturn: def cog_mgr(self) -> NoReturn:

View File

@ -13,14 +13,6 @@ def init_global_checks(bot):
""" """
return ctx.channel.permissions_for(ctx.me).send_messages return ctx.channel.permissions_for(ctx.me).send_messages
@bot.check_once
def actually_up(ctx) -> bool:
"""
Uptime is set during the initial startup process.
If this hasn't been set, we should assume the bot isn't ready yet.
"""
return ctx.bot.uptime is not None
@bot.check_once @bot.check_once
async def whiteblacklist_checks(ctx) -> bool: async def whiteblacklist_checks(ctx) -> bool:
return await ctx.bot.allowed_by_whitelist_blacklist(ctx.author) return await ctx.bot.allowed_by_whitelist_blacklist(ctx.author)