summaryrefslogtreecommitdiffstats
path: root/tqdm/contrib/slack.py
blob: b478d9236ee4fb835285a0f5d673efb57ca1f7bc (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
"""
Sends updates to a Slack app.

Usage:
>>> from tqdm.contrib.slack import tqdm, trange
>>> for i in trange(10, token='{token}', channel='{channel}'):
...     ...

![screenshot](https://img.tqdm.ml/screenshot-slack.png)
"""
from __future__ import absolute_import

import logging
from os import getenv

try:
    from slack_sdk import WebClient
except ImportError:
    raise ImportError("Please `pip install slack-sdk`")

from ..auto import tqdm as tqdm_auto
from ..utils import _range
from .utils_worker import MonoWorker

__author__ = {"github.com/": ["0x2b3bfa0", "casperdcl"]}
__all__ = ['SlackIO', 'tqdm_slack', 'tsrange', 'tqdm', 'trange']


class SlackIO(MonoWorker):
    """Non-blocking file-like IO using a Slack app."""
    def __init__(self, token, channel):
        """Creates a new message in the given `channel`."""
        super(SlackIO, self).__init__()
        self.client = WebClient(token=token)
        self.text = self.__class__.__name__
        try:
            self.message = self.client.chat_postMessage(channel=channel, text=self.text)
        except Exception as e:
            tqdm_auto.write(str(e))
            self.message = None

    def write(self, s):
        """Replaces internal `message`'s text with `s`."""
        if not s:
            s = "..."
        s = s.replace('\r', '').strip()
        if s == self.text:
            return  # skip duplicate message
        message = self.message
        if message is None:
            return
        self.text = s
        try:
            future = self.submit(self.client.chat_update, channel=message['channel'],
                                 ts=message['ts'], text='`' + s + '`')
        except Exception as e:
            tqdm_auto.write(str(e))
        else:
            return future


class tqdm_slack(tqdm_auto):
    """
    Standard `tqdm.auto.tqdm` but also sends updates to a Slack app.
    May take a few seconds to create (`__init__`).

    - create a Slack app with the `chat:write` scope & invite it to a
      channel: <https://api.slack.com/authentication/basics>
    - copy the bot `{token}` & `{channel}` and paste below
    >>> from tqdm.contrib.slack import tqdm, trange
    >>> for i in tqdm(iterable, token='{token}', channel='{channel}'):
    ...     ...
    """
    def __init__(self, *args, **kwargs):
        """
        Parameters
        ----------
        token  : str, required. Slack token
            [default: ${TQDM_SLACK_TOKEN}].
        channel  : int, required. Slack channel
            [default: ${TQDM_SLACK_CHANNEL}].
        mininterval  : float, optional.
          Minimum of [default: 1.5] to avoid rate limit.

        See `tqdm.auto.tqdm.__init__` for other parameters.
        """
        if not kwargs.get('disable'):
            kwargs = kwargs.copy()
            logging.getLogger("HTTPClient").setLevel(logging.WARNING)
            self.sio = SlackIO(
                kwargs.pop('token', getenv("TQDM_SLACK_TOKEN")),
                kwargs.pop('channel', getenv("TQDM_SLACK_CHANNEL")))
            kwargs['mininterval'] = max(1.5, kwargs.get('mininterval', 1.5))
        super(tqdm_slack, self).__init__(*args, **kwargs)

    def display(self, **kwargs):
        super(tqdm_slack, self).display(**kwargs)
        fmt = self.format_dict
        if fmt.get('bar_format', None):
            fmt['bar_format'] = fmt['bar_format'].replace(
                '<bar/>', '`{bar:10}`').replace('{bar}', '`{bar:10u}`')
        else:
            fmt['bar_format'] = '{l_bar}`{bar:10}`{r_bar}'
        if fmt['ascii'] is False:
            fmt['ascii'] = [":black_square:", ":small_blue_diamond:", ":large_blue_diamond:",
                            ":large_blue_square:"]
            fmt['ncols'] = 336
        self.sio.write(self.format_meter(**fmt))

    def clear(self, *args, **kwargs):
        super(tqdm_slack, self).clear(*args, **kwargs)
        if not self.disable:
            self.sio.write("")


def tsrange(*args, **kwargs):
    """
    A shortcut for `tqdm.contrib.slack.tqdm(xrange(*args), **kwargs)`.
    On Python3+, `range` is used instead of `xrange`.
    """
    return tqdm_slack(_range(*args), **kwargs)


# Aliases
tqdm = tqdm_slack
trange = tsrange