discord.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  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. import logging
  10. from os import getenv
  11. try:
  12. from disco.client import Client, ClientConfig
  13. except ImportError:
  14. raise ImportError("Please `pip install disco-py`")
  15. from ..auto import tqdm as tqdm_auto
  16. from .utils_worker import MonoWorker
  17. __author__ = {"github.com/": ["casperdcl"]}
  18. __all__ = ['DiscordIO', 'tqdm_discord', 'tdrange', 'tqdm', 'trange']
  19. class DiscordIO(MonoWorker):
  20. """Non-blocking file-like IO using a Discord Bot."""
  21. def __init__(self, token, channel_id):
  22. """Creates a new message in the given `channel_id`."""
  23. super(DiscordIO, self).__init__()
  24. config = ClientConfig()
  25. config.token = token
  26. client = Client(config)
  27. self.text = self.__class__.__name__
  28. try:
  29. self.message = client.api.channels_messages_create(channel_id, self.text)
  30. except Exception as e:
  31. tqdm_auto.write(str(e))
  32. self.message = None
  33. def write(self, s):
  34. """Replaces internal `message`'s text with `s`."""
  35. if not s:
  36. s = "..."
  37. s = s.replace('\r', '').strip()
  38. if s == self.text:
  39. return # skip duplicate message
  40. message = self.message
  41. if message is None:
  42. return
  43. self.text = s
  44. try:
  45. future = self.submit(message.edit, '`' + s + '`')
  46. except Exception as e:
  47. tqdm_auto.write(str(e))
  48. else:
  49. return future
  50. class tqdm_discord(tqdm_auto):
  51. """
  52. Standard `tqdm.auto.tqdm` but also sends updates to a Discord Bot.
  53. May take a few seconds to create (`__init__`).
  54. - create a discord bot (not public, no requirement of OAuth2 code
  55. grant, only send message permissions) & invite it to a channel:
  56. <https://discordpy.readthedocs.io/en/latest/discord.html>
  57. - copy the bot `{token}` & `{channel_id}` and paste below
  58. >>> from tqdm.contrib.discord import tqdm, trange
  59. >>> for i in tqdm(iterable, token='{token}', channel_id='{channel_id}'):
  60. ... ...
  61. """
  62. def __init__(self, *args, **kwargs):
  63. """
  64. Parameters
  65. ----------
  66. token : str, required. Discord token
  67. [default: ${TQDM_DISCORD_TOKEN}].
  68. channel_id : int, required. Discord channel ID
  69. [default: ${TQDM_DISCORD_CHANNEL_ID}].
  70. mininterval : float, optional.
  71. Minimum of [default: 1.5] to avoid rate limit.
  72. See `tqdm.auto.tqdm.__init__` for other parameters.
  73. """
  74. if not kwargs.get('disable'):
  75. kwargs = kwargs.copy()
  76. logging.getLogger("HTTPClient").setLevel(logging.WARNING)
  77. self.dio = DiscordIO(
  78. kwargs.pop('token', getenv("TQDM_DISCORD_TOKEN")),
  79. kwargs.pop('channel_id', getenv("TQDM_DISCORD_CHANNEL_ID")))
  80. kwargs['mininterval'] = max(1.5, kwargs.get('mininterval', 1.5))
  81. super(tqdm_discord, self).__init__(*args, **kwargs)
  82. def display(self, **kwargs):
  83. super(tqdm_discord, self).display(**kwargs)
  84. fmt = self.format_dict
  85. if fmt.get('bar_format', None):
  86. fmt['bar_format'] = fmt['bar_format'].replace(
  87. '<bar/>', '{bar:10u}').replace('{bar}', '{bar:10u}')
  88. else:
  89. fmt['bar_format'] = '{l_bar}{bar:10u}{r_bar}'
  90. self.dio.write(self.format_meter(**fmt))
  91. def clear(self, *args, **kwargs):
  92. super(tqdm_discord, self).clear(*args, **kwargs)
  93. if not self.disable:
  94. self.dio.write("")
  95. def tdrange(*args, **kwargs):
  96. """Shortcut for `tqdm.contrib.discord.tqdm(range(*args), **kwargs)`."""
  97. return tqdm_discord(range(*args), **kwargs)
  98. # Aliases
  99. tqdm = tqdm_discord
  100. trange = tdrange