telegram.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  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](
  8. https://raw.githubusercontent.com/tqdm/img/src/screenshot-telegram.gif)
  9. """
  10. from __future__ import absolute_import
  11. from os import getenv
  12. from requests import Session
  13. from tqdm.auto import tqdm as tqdm_auto
  14. from tqdm.utils import _range
  15. from .utils_worker import MonoWorker
  16. __author__ = {"github.com/": ["casperdcl"]}
  17. __all__ = ['TelegramIO', 'tqdm_telegram', 'ttgrange', 'tqdm', 'trange']
  18. class TelegramIO(MonoWorker):
  19. """Non-blocking file-like IO using a Telegram Bot."""
  20. API = 'https://api.telegram.org/bot'
  21. def __init__(self, token, chat_id):
  22. """Creates a new message in the given `chat_id`."""
  23. super(TelegramIO, self).__init__()
  24. self.token = token
  25. self.chat_id = chat_id
  26. self.session = session = Session()
  27. self.text = self.__class__.__name__
  28. try:
  29. res = session.post(
  30. self.API + '%s/sendMessage' % self.token,
  31. data=dict(text='`' + self.text + '`', chat_id=self.chat_id,
  32. parse_mode='MarkdownV2'))
  33. except Exception as e:
  34. tqdm_auto.write(str(e))
  35. else:
  36. self.message_id = res.json()['result']['message_id']
  37. def write(self, s):
  38. """Replaces internal `message_id`'s text with `s`."""
  39. if not s:
  40. return
  41. s = s.replace('\r', '').strip()
  42. if s == self.text:
  43. return # avoid duplicate message Bot error
  44. self.text = s
  45. try:
  46. future = self.submit(
  47. self.session.post,
  48. self.API + '%s/editMessageText' % self.token,
  49. data=dict(
  50. text='`' + s + '`', chat_id=self.chat_id,
  51. message_id=self.message_id, parse_mode='MarkdownV2'))
  52. except Exception as e:
  53. tqdm_auto.write(str(e))
  54. else:
  55. return future
  56. class tqdm_telegram(tqdm_auto):
  57. """
  58. Standard `tqdm.auto.tqdm` but also sends updates to a Telegram Bot.
  59. May take a few seconds to create (`__init__`).
  60. - create a bot <https://core.telegram.org/bots#6-botfather>
  61. - copy its `{token}`
  62. - add the bot to a chat and send it a message such as `/start`
  63. - go to <https://api.telegram.org/bot`{token}`/getUpdates> to find out
  64. the `{chat_id}`
  65. - paste the `{token}` & `{chat_id}` below
  66. >>> from tqdm.contrib.telegram import tqdm, trange
  67. >>> for i in tqdm(iterable, token='{token}', chat_id='{chat_id}'):
  68. ... ...
  69. """
  70. def __init__(self, *args, **kwargs):
  71. """
  72. Parameters
  73. ----------
  74. token : str, required. Telegram token
  75. [default: ${TQDM_TELEGRAM_TOKEN}].
  76. chat_id : str, required. Telegram chat ID
  77. [default: ${TQDM_TELEGRAM_CHAT_ID}].
  78. See `tqdm.auto.tqdm.__init__` for other parameters.
  79. """
  80. kwargs = kwargs.copy()
  81. self.tgio = TelegramIO(
  82. kwargs.pop('token', getenv('TQDM_TELEGRAM_TOKEN')),
  83. kwargs.pop('chat_id', getenv('TQDM_TELEGRAM_CHAT_ID')))
  84. super(tqdm_telegram, self).__init__(*args, **kwargs)
  85. def display(self, **kwargs):
  86. super(tqdm_telegram, self).display(**kwargs)
  87. fmt = self.format_dict
  88. if 'bar_format' in fmt and fmt['bar_format']:
  89. fmt['bar_format'] = fmt['bar_format'].replace('<bar/>', '{bar}')
  90. else:
  91. fmt['bar_format'] = '{l_bar}{bar}{r_bar}'
  92. fmt['bar_format'] = fmt['bar_format'].replace('{bar}', '{bar:10u}')
  93. self.tgio.write(self.format_meter(**fmt))
  94. def __new__(cls, *args, **kwargs):
  95. """
  96. Workaround for mixed-class same-stream nested progressbars.
  97. See [#509](https://github.com/tqdm/tqdm/issues/509)
  98. """
  99. with cls.get_lock():
  100. try:
  101. cls._instances = tqdm_auto._instances
  102. except AttributeError:
  103. pass
  104. instance = super(tqdm_telegram, cls).__new__(cls, *args, **kwargs)
  105. with cls.get_lock():
  106. try:
  107. # `tqdm_auto` may have been changed so update
  108. cls._instances.update(tqdm_auto._instances)
  109. except AttributeError:
  110. pass
  111. tqdm_auto._instances = cls._instances
  112. return instance
  113. def ttgrange(*args, **kwargs):
  114. """
  115. A shortcut for `tqdm.contrib.telegram.tqdm(xrange(*args), **kwargs)`.
  116. On Python3+, `range` is used instead of `xrange`.
  117. """
  118. return tqdm_telegram(_range(*args), **kwargs)
  119. # Aliases
  120. tqdm = tqdm_telegram
  121. trange = ttgrange