slack.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. """
  2. Sends updates to a Slack app.
  3. Usage:
  4. >>> from tqdm.contrib.slack import tqdm, trange
  5. >>> for i in trange(10, token='{token}', channel='{channel}'):
  6. ... ...
  7. ![screenshot](https://tqdm.github.io/img/screenshot-slack.png)
  8. """
  9. import logging
  10. from os import getenv
  11. try:
  12. from slack_sdk import WebClient
  13. except ImportError:
  14. raise ImportError("Please `pip install slack-sdk`")
  15. from ..auto import tqdm as tqdm_auto
  16. from .utils_worker import MonoWorker
  17. __author__ = {"github.com/": ["0x2b3bfa0", "casperdcl"]}
  18. __all__ = ['SlackIO', 'tqdm_slack', 'tsrange', 'tqdm', 'trange']
  19. class SlackIO(MonoWorker):
  20. """Non-blocking file-like IO using a Slack app."""
  21. def __init__(self, token, channel):
  22. """Creates a new message in the given `channel`."""
  23. super(SlackIO, self).__init__()
  24. self.client = WebClient(token=token)
  25. self.text = self.__class__.__name__
  26. try:
  27. self.message = self.client.chat_postMessage(channel=channel, text=self.text)
  28. except Exception as e:
  29. tqdm_auto.write(str(e))
  30. self.message = None
  31. def write(self, s):
  32. """Replaces internal `message`'s text with `s`."""
  33. if not s:
  34. s = "..."
  35. s = s.replace('\r', '').strip()
  36. if s == self.text:
  37. return # skip duplicate message
  38. message = self.message
  39. if message is None:
  40. return
  41. self.text = s
  42. try:
  43. future = self.submit(self.client.chat_update, channel=message['channel'],
  44. ts=message['ts'], text='`' + s + '`')
  45. except Exception as e:
  46. tqdm_auto.write(str(e))
  47. else:
  48. return future
  49. class tqdm_slack(tqdm_auto):
  50. """
  51. Standard `tqdm.auto.tqdm` but also sends updates to a Slack app.
  52. May take a few seconds to create (`__init__`).
  53. - create a Slack app with the `chat:write` scope & invite it to a
  54. channel: <https://api.slack.com/authentication/basics>
  55. - copy the bot `{token}` & `{channel}` and paste below
  56. >>> from tqdm.contrib.slack import tqdm, trange
  57. >>> for i in tqdm(iterable, token='{token}', channel='{channel}'):
  58. ... ...
  59. """
  60. def __init__(self, *args, **kwargs):
  61. """
  62. Parameters
  63. ----------
  64. token : str, required. Slack token
  65. [default: ${TQDM_SLACK_TOKEN}].
  66. channel : int, required. Slack channel
  67. [default: ${TQDM_SLACK_CHANNEL}].
  68. mininterval : float, optional.
  69. Minimum of [default: 1.5] to avoid rate limit.
  70. See `tqdm.auto.tqdm.__init__` for other parameters.
  71. """
  72. if not kwargs.get('disable'):
  73. kwargs = kwargs.copy()
  74. logging.getLogger("HTTPClient").setLevel(logging.WARNING)
  75. self.sio = SlackIO(
  76. kwargs.pop('token', getenv("TQDM_SLACK_TOKEN")),
  77. kwargs.pop('channel', getenv("TQDM_SLACK_CHANNEL")))
  78. kwargs['mininterval'] = max(1.5, kwargs.get('mininterval', 1.5))
  79. super(tqdm_slack, self).__init__(*args, **kwargs)
  80. def display(self, **kwargs):
  81. super(tqdm_slack, self).display(**kwargs)
  82. fmt = self.format_dict
  83. if fmt.get('bar_format', None):
  84. fmt['bar_format'] = fmt['bar_format'].replace(
  85. '<bar/>', '`{bar:10}`').replace('{bar}', '`{bar:10u}`')
  86. else:
  87. fmt['bar_format'] = '{l_bar}`{bar:10}`{r_bar}'
  88. if fmt['ascii'] is False:
  89. fmt['ascii'] = [":black_square:", ":small_blue_diamond:", ":large_blue_diamond:",
  90. ":large_blue_square:"]
  91. fmt['ncols'] = 336
  92. self.sio.write(self.format_meter(**fmt))
  93. def clear(self, *args, **kwargs):
  94. super(tqdm_slack, self).clear(*args, **kwargs)
  95. if not self.disable:
  96. self.sio.write("")
  97. def tsrange(*args, **kwargs):
  98. """Shortcut for `tqdm.contrib.slack.tqdm(range(*args), **kwargs)`."""
  99. return tqdm_slack(range(*args), **kwargs)
  100. # Aliases
  101. tqdm = tqdm_slack
  102. trange = tsrange