codekingpro/portable-devtools

1"""2Sends updates to a Discord bot.3 4Usage:5>>> from tqdm.contrib.discord import tqdm, trange6>>> for i in trange(10, token='{token}', channel_id='{channel_id}'):7...     ...8 9![screenshot](https://tqdm.github.io/img/screenshot-discord.png)10"""11from os import getenv12from warnings import warn13 14from requests import Session15from requests.utils import default_user_agent16 17from ..auto import tqdm as tqdm_auto18from ..std import TqdmWarning19from ..version import __version__20from .utils_worker import MonoWorker21 22__author__ = {"github.com/": ["casperdcl", "guigoruiz1"]}23__all__ = ['DiscordIO', 'tqdm_discord', 'tdrange', 'tqdm', 'trange']24 25 26class DiscordIO(MonoWorker):27    """Non-blocking file-like IO using a Discord Bot."""28    API = "https://discord.com/api/v10"29    UA = f"tqdm (https://tqdm.github.io, {__version__}) {default_user_agent()}"30 31    def __init__(self, token, channel_id):32        """Creates a new message in the given `channel_id`."""33        super().__init__()34        self.token = token35        self.channel_id = channel_id36        self.session = Session()37        self.text = self.__class__.__name__38        self.message_id39 40    @property41    def message_id(self):42        if hasattr(self, '_message_id'):43            return self._message_id44        try:45            res = self.session.post(46                f'{self.API}/channels/{self.channel_id}/messages',47                headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA},48                json={'content': f"`{self.text}`"}).json()49        except Exception as e:50            tqdm_auto.write(str(e))51        else:52            if res.get('error_code') == 429:53                warn("Creation rate limit: try increasing `mininterval`.",54                     TqdmWarning, stacklevel=2)55            else:56                self._message_id = res['id']57                return self._message_id58 59    def write(self, s):60        """Replaces internal `message_id`'s text with `s`."""61        if not s:62            s = "..."63        s = s.replace('\r', '').strip()64        if s == self.text:65            return  # avoid duplicate message Bot error66        message_id = self.message_id67        if message_id is None:68            return69        self.text = s70        try:71            future = self.submit(72                self.session.patch,73                f'{self.API}/channels/{self.channel_id}/messages/{message_id}',74                headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA},75                json={'content': f"`{self.text}`"})76        except Exception as e:77            tqdm_auto.write(str(e))78        else:79            return future80 81    def delete(self):82        """Deletes internal `message_id`."""83        try:84            future = self.submit(85                self.session.delete,86                f'{self.API}/channels/{self.channel_id}/messages/{self.message_id}',87                headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA})88        except Exception as e:89            tqdm_auto.write(str(e))90        else:91            return future92 93 94class tqdm_discord(tqdm_auto):95    """96    Standard `tqdm.auto.tqdm` but also sends updates to a Discord Bot.97    May take a few seconds to create (`__init__`).98 99    - create a discord bot (not public, no requirement of OAuth2 code100      grant, only send message permissions) & invite it to a channel:101      <https://discordpy.readthedocs.io/en/latest/discord.html>102    - copy the bot `{token}` & `{channel_id}` and paste below103 104    >>> from tqdm.contrib.discord import tqdm, trange105    >>> for i in tqdm(iterable, token='{token}', channel_id='{channel_id}'):106    ...     ...107    """108    def __init__(self, *args, **kwargs):109        """110        Parameters111        ----------112        token  : str, required. Discord bot token113            [default: ${TQDM_DISCORD_TOKEN}].114        channel_id  : int, required. Discord channel ID115            [default: ${TQDM_DISCORD_CHANNEL_ID}].116 117        See `tqdm.auto.tqdm.__init__` for other parameters.118        """119        if not kwargs.get('disable'):120            kwargs = kwargs.copy()121            self.dio = DiscordIO(122                kwargs.pop('token', getenv('TQDM_DISCORD_TOKEN')),123                kwargs.pop('channel_id', getenv('TQDM_DISCORD_CHANNEL_ID')))124        super().__init__(*args, **kwargs)125 126    def display(self, **kwargs):127        super().display(**kwargs)128        fmt = self.format_dict129        if fmt.get('bar_format', None):130            fmt['bar_format'] = fmt['bar_format'].replace(131                '<bar/>', '{bar:10u}').replace('{bar}', '{bar:10u}')132        else:133            fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}'134        self.dio.write(self.format_meter(**fmt))135 136    def clear(self, *args, **kwargs):137        super().clear(*args, **kwargs)138        if not self.disable:139            self.dio.write("")140 141    def close(self):142        if self.disable:143            return144        super().close()145        if not (self.leave or (self.leave is None and self.pos == 0)):146            self.dio.delete()147 148 149def tdrange(*args, **kwargs):150    """Shortcut for `tqdm.contrib.discord.tqdm(range(*args), **kwargs)`."""151    return tqdm_discord(range(*args), **kwargs)152 153 154# Aliases155tqdm = tqdm_discord156trange = tdrange157