1# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
2#
3# Permission to use, copy, modify, and distribute this software for any
4# purpose with or without fee is hereby granted, provided that the above
5# copyright notice and this permission notice appear in all copies.
6#
7# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
8# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
9# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
10# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
11# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
12# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
13# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
14
15from __future__ import division, print_function
16
17from collections import deque
18from datetime import timedelta
19from math import ceil
20from sys import stderr
21try:
22    from time import monotonic
23except ImportError:
24    from time import time as monotonic
25
26
27__version__ = '1.5'
28
29HIDE_CURSOR = '\x1b[?25l'
30SHOW_CURSOR = '\x1b[?25h'
31
32
33class Infinite(object):
34    file = stderr
35    sma_window = 10         # Simple Moving Average window
36    check_tty = True
37    hide_cursor = True
38
39    def __init__(self, message='', **kwargs):
40        self.index = 0
41        self.start_ts = monotonic()
42        self.avg = 0
43        self._avg_update_ts = self.start_ts
44        self._ts = self.start_ts
45        self._xput = deque(maxlen=self.sma_window)
46        for key, val in kwargs.items():
47            setattr(self, key, val)
48
49        self._width = 0
50        self.message = message
51
52        if self.file and self.is_tty():
53            if self.hide_cursor:
54                print(HIDE_CURSOR, end='', file=self.file)
55            print(self.message, end='', file=self.file)
56            self.file.flush()
57
58    def __getitem__(self, key):
59        if key.startswith('_'):
60            return None
61        return getattr(self, key, None)
62
63    @property
64    def elapsed(self):
65        return int(monotonic() - self.start_ts)
66
67    @property
68    def elapsed_td(self):
69        return timedelta(seconds=self.elapsed)
70
71    def update_avg(self, n, dt):
72        if n > 0:
73            xput_len = len(self._xput)
74            self._xput.append(dt / n)
75            now = monotonic()
76            # update when we're still filling _xput, then after every second
77            if (xput_len < self.sma_window or
78                    now - self._avg_update_ts > 1):
79                self.avg = sum(self._xput) / len(self._xput)
80                self._avg_update_ts = now
81
82    def update(self):
83        pass
84
85    def start(self):
86        pass
87
88    def clearln(self):
89        if self.file and self.is_tty():
90            print('\r\x1b[K', end='', file=self.file)
91
92    def write(self, s):
93        if self.file and self.is_tty():
94            line = self.message + s.ljust(self._width)
95            print('\r' + line, end='', file=self.file)
96            self._width = max(self._width, len(s))
97            self.file.flush()
98
99    def writeln(self, line):
100        if self.file and self.is_tty():
101            self.clearln()
102            print(line, end='', file=self.file)
103            self.file.flush()
104
105    def finish(self):
106        if self.file and self.is_tty():
107            print(file=self.file)
108            if self.hide_cursor:
109                print(SHOW_CURSOR, end='', file=self.file)
110
111    def is_tty(self):
112        return self.file.isatty() if self.check_tty else True
113
114    def next(self, n=1):
115        now = monotonic()
116        dt = now - self._ts
117        self.update_avg(n, dt)
118        self._ts = now
119        self.index = self.index + n
120        self.update()
121
122    def iter(self, it):
123        with self:
124            for x in it:
125                yield x
126                self.next()
127
128    def __enter__(self):
129        self.start()
130        return self
131
132    def __exit__(self, exc_type, exc_val, exc_tb):
133        self.finish()
134
135
136class Progress(Infinite):
137    def __init__(self, *args, **kwargs):
138        super(Progress, self).__init__(*args, **kwargs)
139        self.max = kwargs.get('max', 100)
140
141    @property
142    def eta(self):
143        return int(ceil(self.avg * self.remaining))
144
145    @property
146    def eta_td(self):
147        return timedelta(seconds=self.eta)
148
149    @property
150    def percent(self):
151        return self.progress * 100
152
153    @property
154    def progress(self):
155        return min(1, self.index / self.max)
156
157    @property
158    def remaining(self):
159        return max(self.max - self.index, 0)
160
161    def start(self):
162        self.update()
163
164    def goto(self, index):
165        incr = index - self.index
166        self.next(incr)
167
168    def iter(self, it):
169        try:
170            self.max = len(it)
171        except TypeError:
172            pass
173
174        with self:
175            for x in it:
176                yield x
177                self.next()
178