telegram.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. """
  2. Sends updates to a Telegram bot.
  3. Usage:
  4. >>> from tqdm.contrib.telegram import tqdm, trange
  5. >>> for i in trange(10, token='{token}', chat_id='{chat_id}'):
  6. ... ...
  7. ![screenshot](https://tqdm.github.io/img/screenshot-telegram.gif)
  8. """
  9. from os import getenv
  10. from warnings import warn
  11. from requests import Session
  12. from ..auto import tqdm as tqdm_auto
  13. from ..std import TqdmWarning
  14. from .utils_worker import MonoWorker
  15. __author__ = {"github.com/": ["casperdcl"]}
  16. __all__ = ['TelegramIO', 'tqdm_telegram', 'ttgrange', 'tqdm', 'trange']
  17. class TelegramIO(MonoWorker):
  18. """Non-blocking file-like IO using a Telegram Bot."""
  19. API = 'https://api.telegram.org/bot'
  20. def __init__(self, token, chat_id):
  21. """Creates a new message in the given `chat_id`."""
  22. super().__init__()
  23. self.token = token
  24. self.chat_id = chat_id
  25. self.session = Session()
  26. self.text = self.__class__.__name__
  27. self.message_id
  28. @property
  29. def message_id(self):
  30. if hasattr(self, '_message_id'):
  31. return self._message_id
  32. try:
  33. res = self.session.post(
  34. self.API + '%s/sendMessage' % self.token,
  35. data={'text': '`' + self.text + '`', 'chat_id': self.chat_id,
  36. 'parse_mode': 'MarkdownV2'}).json()
  37. except Exception as e:
  38. tqdm_auto.write(str(e))
  39. else:
  40. if res.get('error_code') == 429:
  41. warn("Creation rate limit: try increasing `mininterval`.",
  42. TqdmWarning, stacklevel=2)
  43. else:
  44. self._message_id = res['result']['message_id']
  45. return self._message_id
  46. def write(self, s):
  47. """Replaces internal `message_id`'s text with `s`."""
  48. if not s:
  49. s = "..."
  50. s = s.replace('\r', '').strip()
  51. if s == self.text:
  52. return # avoid duplicate message Bot error
  53. message_id = self.message_id
  54. if message_id is None:
  55. return
  56. self.text = s
  57. try:
  58. future = self.submit(
  59. self.session.post, self.API + '%s/editMessageText' % self.token,
  60. data={'text': '`' + s + '`', 'chat_id': self.chat_id,
  61. 'message_id': message_id, 'parse_mode': 'MarkdownV2'})
  62. except Exception as e:
  63. tqdm_auto.write(str(e))
  64. else:
  65. return future
  66. def delete(self):
  67. """Deletes internal `message_id`."""
  68. try:
  69. future = self.submit(
  70. self.session.post, self.API + '%s/deleteMessage' % self.token,
  71. data={'chat_id': self.chat_id, 'message_id': self.message_id})
  72. except Exception as e:
  73. tqdm_auto.write(str(e))
  74. else:
  75. return future
  76. class tqdm_telegram(tqdm_auto):
  77. """
  78. Standard `tqdm.auto.tqdm` but also sends updates to a Telegram Bot.
  79. May take a few seconds to create (`__init__`).
  80. - create a bot <https://core.telegram.org/bots#6-botfather>
  81. - copy its `{token}`
  82. - add the bot to a chat and send it a message such as `/start`
  83. - go to <https://api.telegram.org/bot`{token}`/getUpdates> to find out
  84. the `{chat_id}`
  85. - paste the `{token}` & `{chat_id}` below
  86. >>> from tqdm.contrib.telegram import tqdm, trange
  87. >>> for i in tqdm(iterable, token='{token}', chat_id='{chat_id}'):
  88. ... ...
  89. """
  90. def __init__(self, *args, **kwargs):
  91. """
  92. Parameters
  93. ----------
  94. token : str, required. Telegram token
  95. [default: ${TQDM_TELEGRAM_TOKEN}].
  96. chat_id : str, required. Telegram chat ID
  97. [default: ${TQDM_TELEGRAM_CHAT_ID}].
  98. See `tqdm.auto.tqdm.__init__` for other parameters.
  99. """
  100. if not kwargs.get('disable'):
  101. kwargs = kwargs.copy()
  102. self.tgio = TelegramIO(
  103. kwargs.pop('token', getenv('TQDM_TELEGRAM_TOKEN')),
  104. kwargs.pop('chat_id', getenv('TQDM_TELEGRAM_CHAT_ID')))
  105. super().__init__(*args, **kwargs)
  106. def display(self, **kwargs):
  107. super().display(**kwargs)
  108. fmt = self.format_dict
  109. if fmt.get('bar_format', None):
  110. fmt['bar_format'] = fmt['bar_format'].replace(
  111. '<bar/>', '{bar:10u}').replace('{bar}', '{bar:10u}')
  112. else:
  113. fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}'
  114. self.tgio.write(self.format_meter(**fmt))
  115. def clear(self, *args, **kwargs):
  116. super().clear(*args, **kwargs)
  117. if not self.disable:
  118. self.tgio.write("")
  119. def close(self):
  120. if self.disable:
  121. return
  122. super().close()
  123. if not (self.leave or (self.leave is None and self.pos == 0)):
  124. self.tgio.delete()
  125. def ttgrange(*args, **kwargs):
  126. """Shortcut for `tqdm.contrib.telegram.tqdm(range(*args), **kwargs)`."""
  127. return tqdm_telegram(range(*args), **kwargs)
  128. # Aliases
  129. tqdm = tqdm_telegram
  130. trange = ttgrange