diff --git a/pyproject.toml b/pyproject.toml index ec1a63e4..5651e88e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,7 @@ dependencies = ['colorama; platform_system == "Windows"'] [project.optional-dependencies] dev = ["pytest>=6", "pytest-cov", "pytest-timeout", "pytest-xdist"] +discord = ["requests"] slack = ["slack-sdk"] telegram = ["requests"] notebook = ["ipywidgets>=6"] diff --git a/tqdm/contrib/discord.py b/tqdm/contrib/discord.py index fd663f98..574baa84 100644 --- a/tqdm/contrib/discord.py +++ b/tqdm/contrib/discord.py @@ -8,49 +8,83 @@ Usage: ![screenshot](https://tqdm.github.io/img/screenshot-discord.png) """ -import logging from os import getenv +from warnings import warn -try: - from disco.client import Client, ClientConfig -except ImportError: - raise ImportError("Please `pip install disco-py`") +from requests import Session +from requests.utils import default_user_agent from ..auto import tqdm as tqdm_auto +from ..std import TqdmWarning +from ..version import __version__ from .utils_worker import MonoWorker -__author__ = {"github.com/": ["casperdcl"]} +__author__ = {"github.com/": ["casperdcl", "guigoruiz1"]} __all__ = ['DiscordIO', 'tqdm_discord', 'tdrange', 'tqdm', 'trange'] class DiscordIO(MonoWorker): """Non-blocking file-like IO using a Discord Bot.""" + API = "https://discord.com/api/v10" + UA = f"tqdm (https://tqdm.github.io, {__version__}) {default_user_agent()}" + def __init__(self, token, channel_id): """Creates a new message in the given `channel_id`.""" super().__init__() - config = ClientConfig() - config.token = token - client = Client(config) + self.token = token + self.channel_id = channel_id + self.session = Session() self.text = self.__class__.__name__ + self.message_id + + @property + def message_id(self): + if hasattr(self, '_message_id'): + return self._message_id try: - self.message = client.api.channels_messages_create(channel_id, self.text) + res = self.session.post( + f'{self.API}/channels/{self.channel_id}/messages', + headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA}, + json={'content': f"`{self.text}`"}).json() except Exception as e: tqdm_auto.write(str(e)) - self.message = None + else: + if res.get('error_code') == 429: + warn("Creation rate limit: try increasing `mininterval`.", + TqdmWarning, stacklevel=2) + else: + self._message_id = res['id'] + return self._message_id def write(self, s): - """Replaces internal `message`'s text with `s`.""" + """Replaces internal `message_id`'s text with `s`.""" if not s: s = "..." s = s.replace('\r', '').strip() if s == self.text: - return # skip duplicate message - message = self.message - if message is None: + return # avoid duplicate message Bot error + message_id = self.message_id + if message_id is None: return self.text = s try: - future = self.submit(message.edit, '`' + s + '`') + future = self.submit( + self.session.patch, + f'{self.API}/channels/{self.channel_id}/messages/{message_id}', + headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA}, + json={'content': f"`{self.text}`"}) + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + def delete(self): + """Deletes internal `message_id`.""" + try: + future = self.submit( + self.session.delete, + f'{self.API}/channels/{self.channel_id}/messages/{self.message_id}', + headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA}) except Exception as e: tqdm_auto.write(str(e)) else: @@ -75,22 +109,18 @@ class tqdm_discord(tqdm_auto): """ Parameters ---------- - token : str, required. Discord token + token : str, required. Discord bot token [default: ${TQDM_DISCORD_TOKEN}]. channel_id : int, required. Discord channel ID [default: ${TQDM_DISCORD_CHANNEL_ID}]. - mininterval : float, optional. - Minimum of [default: 1.5] to avoid rate limit. See `tqdm.auto.tqdm.__init__` for other parameters. """ if not kwargs.get('disable'): kwargs = kwargs.copy() - logging.getLogger("HTTPClient").setLevel(logging.WARNING) self.dio = DiscordIO( - kwargs.pop('token', getenv("TQDM_DISCORD_TOKEN")), - kwargs.pop('channel_id', getenv("TQDM_DISCORD_CHANNEL_ID"))) - kwargs['mininterval'] = max(1.5, kwargs.get('mininterval', 1.5)) + kwargs.pop('token', getenv('TQDM_DISCORD_TOKEN')), + kwargs.pop('channel_id', getenv('TQDM_DISCORD_CHANNEL_ID'))) super().__init__(*args, **kwargs) def display(self, **kwargs): @@ -108,6 +138,13 @@ class tqdm_discord(tqdm_auto): if not self.disable: self.dio.write("") + def close(self): + if self.disable: + return + super().close() + if not (self.leave or (self.leave is None and self.pos == 0)): + self.dio.delete() + def tdrange(*args, **kwargs): """Shortcut for `tqdm.contrib.discord.tqdm(range(*args), **kwargs)`."""