discord.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156
  1. """
  2. Sends updates to a Discord bot.
  3. Usage:
  4. >>> from tqdm.contrib.discord import tqdm, trange
  5. >>> for i in trange(10, token='{token}', channel_id='{channel_id}'):
  6. ... ...
  7. ![screenshot](https://tqdm.github.io/img/screenshot-discord.png)
  8. """
  9. from os import getenv
  10. from warnings import warn
  11. from requests import Session
  12. from requests.utils import default_user_agent
  13. from ..auto import tqdm as tqdm_auto
  14. from ..std import TqdmWarning
  15. from ..version import __version__
  16. from .utils_worker import MonoWorker
  17. __author__ = {"github.com/": ["casperdcl", "guigoruiz1"]}
  18. __all__ = ['DiscordIO', 'tqdm_discord', 'tdrange', 'tqdm', 'trange']
  19. class DiscordIO(MonoWorker):
  20. """Non-blocking file-like IO using a Discord Bot."""
  21. API = "https://discord.com/api/v10"
  22. UA = f"tqdm (https://tqdm.github.io, {__version__}) {default_user_agent()}"
  23. def __init__(self, token, channel_id):
  24. """Creates a new message in the given `channel_id`."""
  25. super().__init__()
  26. self.token = token
  27. self.channel_id = channel_id
  28. self.session = Session()
  29. self.text = self.__class__.__name__
  30. self.message_id
  31. @property
  32. def message_id(self):
  33. if hasattr(self, '_message_id'):
  34. return self._message_id
  35. try:
  36. res = self.session.post(
  37. f'{self.API}/channels/{self.channel_id}/messages',
  38. headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA},
  39. json={'content': f"`{self.text}`"}).json()
  40. except Exception as e:
  41. tqdm_auto.write(str(e))
  42. else:
  43. if res.get('error_code') == 429:
  44. warn("Creation rate limit: try increasing `mininterval`.",
  45. TqdmWarning, stacklevel=2)
  46. else:
  47. self._message_id = res['id']
  48. return self._message_id
  49. def write(self, s):
  50. """Replaces internal `message_id`'s text with `s`."""
  51. if not s:
  52. s = "..."
  53. s = s.replace('\r', '').strip()
  54. if s == self.text:
  55. return # avoid duplicate message Bot error
  56. message_id = self.message_id
  57. if message_id is None:
  58. return
  59. self.text = s
  60. try:
  61. future = self.submit(
  62. self.session.patch,
  63. f'{self.API}/channels/{self.channel_id}/messages/{message_id}',
  64. headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA},
  65. json={'content': f"`{self.text}`"})
  66. except Exception as e:
  67. tqdm_auto.write(str(e))
  68. else:
  69. return future
  70. def delete(self):
  71. """Deletes internal `message_id`."""
  72. try:
  73. future = self.submit(
  74. self.session.delete,
  75. f'{self.API}/channels/{self.channel_id}/messages/{self.message_id}',
  76. headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA})
  77. except Exception as e:
  78. tqdm_auto.write(str(e))
  79. else:
  80. return future
  81. class tqdm_discord(tqdm_auto):
  82. """
  83. Standard `tqdm.auto.tqdm` but also sends updates to a Discord Bot.
  84. May take a few seconds to create (`__init__`).
  85. - create a discord bot (not public, no requirement of OAuth2 code
  86. grant, only send message permissions) & invite it to a channel:
  87. <https://discordpy.readthedocs.io/en/latest/discord.html>
  88. - copy the bot `{token}` & `{channel_id}` and paste below
  89. >>> from tqdm.contrib.discord import tqdm, trange
  90. >>> for i in tqdm(iterable, token='{token}', channel_id='{channel_id}'):
  91. ... ...
  92. """
  93. def __init__(self, *args, **kwargs):
  94. """
  95. Parameters
  96. ----------
  97. token : str, required. Discord bot token
  98. [default: ${TQDM_DISCORD_TOKEN}].
  99. channel_id : int, required. Discord channel ID
  100. [default: ${TQDM_DISCORD_CHANNEL_ID}].
  101. See `tqdm.auto.tqdm.__init__` for other parameters.
  102. """
  103. if not kwargs.get('disable'):
  104. kwargs = kwargs.copy()
  105. self.dio = DiscordIO(
  106. kwargs.pop('token', getenv('TQDM_DISCORD_TOKEN')),
  107. kwargs.pop('channel_id', getenv('TQDM_DISCORD_CHANNEL_ID')))
  108. super().__init__(*args, **kwargs)
  109. def display(self, **kwargs):
  110. super().display(**kwargs)
  111. fmt = self.format_dict
  112. if fmt.get('bar_format', None):
  113. fmt['bar_format'] = fmt['bar_format'].replace(
  114. '<bar/>', '{bar:10u}').replace('{bar}', '{bar:10u}')
  115. else:
  116. fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}'
  117. self.dio.write(self.format_meter(**fmt))
  118. def clear(self, *args, **kwargs):
  119. super().clear(*args, **kwargs)
  120. if not self.disable:
  121. self.dio.write("")
  122. def close(self):
  123. if self.disable:
  124. return
  125. super().close()
  126. if not (self.leave or (self.leave is None and self.pos == 0)):
  127. self.dio.delete()
  128. def tdrange(*args, **kwargs):
  129. """Shortcut for `tqdm.contrib.discord.tqdm(range(*args), **kwargs)`."""
  130. return tqdm_discord(range(*args), **kwargs)
  131. # Aliases
  132. tqdm = tqdm_discord
  133. trange = tdrange