From b78efea5903a5f511a3bb22f8de32340ca4b954b Mon Sep 17 00:00:00 2001 From: Guilherme Ruiz <57478888+guigoruiz1@users.noreply.github.com> Date: Sun, 26 Nov 2023 15:45:57 +0000 Subject: [PATCH] Using REST API directly --- tqdm/contrib/discord.py | 131 ++++++++++++++++++++++++++-------------- 1 file changed, 85 insertions(+), 46 deletions(-) diff --git a/tqdm/contrib/discord.py b/tqdm/contrib/discord.py index fd663f98..d4572e67 100644 --- a/tqdm/contrib/discord.py +++ b/tqdm/contrib/discord.py @@ -1,56 +1,92 @@ -""" -Sends updates to a Discord bot. - -Usage: ->>> from tqdm.contrib.discord import tqdm, trange ->>> for i in trange(10, token='{token}', channel_id='{channel_id}'): -... ... - -![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`") +import requests from ..auto import tqdm as tqdm_auto +from ..std import TqdmWarning from .utils_worker import MonoWorker +from ..version import __version__ __author__ = {"github.com/": ["casperdcl"]} -__all__ = ['DiscordIO', 'tqdm_discord', 'tdrange', 'tqdm', 'trange'] +__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" + 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.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) + headers = { + "Authorization": f"Bot {self.token}", + "User-Agent": f"TQDM Discord progress bar (http://github.com/tqdm/tqdm, {__version__})", + } + data = {"content": "`" + self.text + "`"} + res = requests.post( + f"{self.API}/channels/{self.channel_id}/messages", + headers=headers, + json=data, + ).json() except Exception as e: tqdm_auto.write(str(e)) - self.message = None + else: + if "id" in res: + 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() + 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 + '`') + headers = { + "Authorization": f"Bot {self.token}", + "User-Agent": f"TQDM Discord progress bar (http://github.com/tqdm/tqdm, {__version__})", + } + data = {"content": "`" + s + "`"} + future = self.submit( + requests.patch, + f"{self.API}/channels/{self.channel_id}/messages/{message_id}", + headers=headers, + json=data, + ) + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + def delete(self): + """Deletes internal `message_id`.""" + try: + headers = { + "Authorization": f"Bot {self.token}", + "User-Agent": f"TQDM Discord progress bar (http://github.com/tqdm/tqdm, {__version__})", + } + future = self.submit( + requests.delete, + f"{self.API}/channels/{self.channel_id}/messages/{self.message_id}", + headers=headers, + ) except Exception as e: tqdm_auto.write(str(e)) else: @@ -62,45 +98,41 @@ class tqdm_discord(tqdm_auto): Standard `tqdm.auto.tqdm` but also sends updates to a Discord Bot. May take a few seconds to create (`__init__`). - - create a discord bot (not public, no requirement of OAuth2 code - grant, only send message permissions) & invite it to a channel: - - - copy the bot `{token}` & `{channel_id}` and paste below - - >>> from tqdm.contrib.discord import tqdm, trange + >>> from tqdm.contrib.discord import tqdm, drange >>> for i in tqdm(iterable, token='{token}', channel_id='{channel_id}'): ... ... """ + def __init__(self, *args, **kwargs): """ Parameters ---------- - token : str, required. Discord token + token : str, required. Discord bot token [default: ${TQDM_DISCORD_TOKEN}]. - channel_id : int, required. Discord channel ID + channel_id : str, 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'): + 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): super().display(**kwargs) fmt = self.format_dict - if fmt.get('bar_format', None): - fmt['bar_format'] = fmt['bar_format'].replace( - '', '{bar:10u}').replace('{bar}', '{bar:10u}') + if fmt.get("bar_format", None): + fmt["bar_format"] = ( + fmt["bar_format"] + .replace("", "{bar:10u}") + .replace("{bar}", "{bar:10u}") + ) else: - fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}' + fmt["bar_format"] = "{l_bar}{bar:10u}{r_bar}" self.dio.write(self.format_meter(**fmt)) def clear(self, *args, **kwargs): @@ -108,6 +140,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)`."""