discord.py ratelimit monitoring.

rl_handler.py

#   Copyright 2022 Michael Hall
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

from importlib.metadata import version
import warnings
from typing import Callable
import logging

d_version = version("discord.py")


class RLTrack(logging.Handler):

    def __init__(self, callback: Callable):
        super().__init__(logging.debug)
        self._callback = callback  # TODO: proper typing of each variant this will be called with.

    log_types = {
        'WebSocket in shard ID %s is ratelimited, waiting %.2f seconds': "WEBSOCKET_RATELIMITED",
        'A rate limit bucket (%s) has been exhausted. Pre-emptively rate limiting...': "HTTP_PRE_RATELIMIT",
        '%s %s received a 429 despite having %s remaining requests. This is a sub-ratelimit.': "HTTP_SUB_RATELIMIT",
        'We are being rate limited. %s %s responded with 429. Timeout of %.2f was too long, erroring instead.': "HTTP_RATELIMIT",
        'We are being rate limited. %s %s responded with 429. Retrying in %.2f seconds.': "HTTP_RATELIMIT",
        'Global rate limit has been hit. Retrying in %.2f seconds.': "GLOBAL_RATELIMIT",
        'Webhook ID %s has exhausted its rate limit bucket (retry: %s).': "WEBHOOK_RATELIMIT",
        'Webhook ID %s is rate limited. Retrying in %.2f seconds.': "WEBHOOK_RATELIMIT",
    }

    def emit(self, record: logging.LogRecord):
        try:
            rl_type = self.log_types[record.msg]
            self._callback(rl_type, *record.args)
        except KeyError:
            pass  # nosec

    def install(self):
        if d_version not in ("2.0.0", "2.0.1", "2.1.0"):
            warnings.warn(f"Unexpected discord.py version {d_version}. This may or may not function as expected.")
        log = logging.getLogger("discord")
        log.addHandler(self)