diff options
Diffstat (limited to 'third_party/python/tqdm/tqdm/contrib/telegram.py')
| -rw-r--r-- | third_party/python/tqdm/tqdm/contrib/telegram.py | 159 |
1 files changed, 159 insertions, 0 deletions
diff --git a/third_party/python/tqdm/tqdm/contrib/telegram.py b/third_party/python/tqdm/tqdm/contrib/telegram.py new file mode 100644 index 0000000000..99cbe8c888 --- /dev/null +++ b/third_party/python/tqdm/tqdm/contrib/telegram.py @@ -0,0 +1,159 @@ +""" +Sends updates to a Telegram bot. + +Usage: +>>> from tqdm.contrib.telegram import tqdm, trange +>>> for i in trange(10, token='{token}', chat_id='{chat_id}'): +... ... + + +""" +from __future__ import absolute_import + +from os import getenv +from warnings import warn + +from requests import Session + +from ..auto import tqdm as tqdm_auto +from ..std import TqdmWarning +from ..utils import _range +from .utils_worker import MonoWorker + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['TelegramIO', 'tqdm_telegram', 'ttgrange', 'tqdm', 'trange'] + + +class TelegramIO(MonoWorker): + """Non-blocking file-like IO using a Telegram Bot.""" + API = 'https://api.telegram.org/bot' + + def __init__(self, token, chat_id): + """Creates a new message in the given `chat_id`.""" + super(TelegramIO, self).__init__() + self.token = token + self.chat_id = chat_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: + res = self.session.post( + self.API + '%s/sendMessage' % self.token, + data={'text': '`' + self.text + '`', 'chat_id': self.chat_id, + 'parse_mode': 'MarkdownV2'}).json() + except Exception as e: + tqdm_auto.write(str(e)) + else: + if res.get('error_code') == 429: + warn("Creation rate limit: try increasing `mininterval`.", + TqdmWarning, stacklevel=2) + else: + self._message_id = res['result']['message_id'] + return self._message_id + + def write(self, s): + """Replaces internal `message_id`'s text with `s`.""" + if not s: + s = "..." + s = s.replace('\r', '').strip() + if s == self.text: + return # avoid duplicate message Bot error + message_id = self.message_id + if message_id is None: + return + self.text = s + try: + future = self.submit( + self.session.post, self.API + '%s/editMessageText' % self.token, + data={'text': '`' + s + '`', 'chat_id': self.chat_id, + 'message_id': message_id, 'parse_mode': 'MarkdownV2'}) + 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.post, self.API + '%s/deleteMessage' % self.token, + data={'chat_id': self.chat_id, 'message_id': self.message_id}) + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + +class tqdm_telegram(tqdm_auto): + """ + Standard `tqdm.auto.tqdm` but also sends updates to a Telegram Bot. + May take a few seconds to create (`__init__`). + + - create a bot <https://core.telegram.org/bots#6-botfather> + - copy its `{token}` + - add the bot to a chat and send it a message such as `/start` + - go to <https://api.telegram.org/bot`{token}`/getUpdates> to find out + the `{chat_id}` + - paste the `{token}` & `{chat_id}` below + + >>> from tqdm.contrib.telegram import tqdm, trange + >>> for i in tqdm(iterable, token='{token}', chat_id='{chat_id}'): + ... ... + """ + def __init__(self, *args, **kwargs): + """ + Parameters + ---------- + token : str, required. Telegram token + [default: ${TQDM_TELEGRAM_TOKEN}]. + chat_id : str, required. Telegram chat ID + [default: ${TQDM_TELEGRAM_CHAT_ID}]. + + See `tqdm.auto.tqdm.__init__` for other parameters. + """ + if not kwargs.get('disable'): + kwargs = kwargs.copy() + self.tgio = TelegramIO( + kwargs.pop('token', getenv('TQDM_TELEGRAM_TOKEN')), + kwargs.pop('chat_id', getenv('TQDM_TELEGRAM_CHAT_ID'))) + super(tqdm_telegram, self).__init__(*args, **kwargs) + + def display(self, **kwargs): + super(tqdm_telegram, self).display(**kwargs) + fmt = self.format_dict + if fmt.get('bar_format', None): + fmt['bar_format'] = fmt['bar_format'].replace( + '<bar/>', '{bar:10u}').replace('{bar}', '{bar:10u}') + else: + fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}' + self.tgio.write(self.format_meter(**fmt)) + + def clear(self, *args, **kwargs): + super(tqdm_telegram, self).clear(*args, **kwargs) + if not self.disable: + self.tgio.write("") + + def close(self): + if self.disable: + return + super(tqdm_telegram, self).close() + if not (self.leave or (self.leave is None and self.pos == 0)): + self.tgio.delete() + + +def ttgrange(*args, **kwargs): + """ + A shortcut for `tqdm.contrib.telegram.tqdm(xrange(*args), **kwargs)`. + On Python3+, `range` is used instead of `xrange`. + """ + return tqdm_telegram(_range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_telegram +trange = ttgrange |