~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/progress.py

merge bzr.dev

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2005 Aaron Bentley <aaron.bentley@utoronto.ca>
 
2
# Copyright (C) 2005, 2006 Canonical <canonical.com>
 
3
#
 
4
#    This program is free software; you can redistribute it and/or modify
 
5
#    it under the terms of the GNU General Public License as published by
 
6
#    the Free Software Foundation; either version 2 of the License, or
 
7
#    (at your option) any later version.
 
8
#
 
9
#    This program is distributed in the hope that it will be useful,
 
10
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
 
11
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
12
#    GNU General Public License for more details.
 
13
#
 
14
#    You should have received a copy of the GNU General Public License
 
15
#    along with this program; if not, write to the Free Software
 
16
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
17
 
 
18
 
 
19
"""Simple text-mode progress indicator.
 
20
 
 
21
To display an indicator, create a ProgressBar object.  Call it,
 
22
passing Progress objects indicating the current state.  When done,
 
23
call clear().
 
24
 
 
25
Progress is suppressed when output is not sent to a terminal, so as
 
26
not to clutter log files.
 
27
"""
 
28
 
 
29
# TODO: should be a global option e.g. --silent that disables progress
 
30
# indicators, preferably without needing to adjust all code that
 
31
# potentially calls them.
 
32
 
 
33
# TODO: If not on a tty perhaps just print '......' for the benefit of IDEs, etc
 
34
 
 
35
# TODO: Optionally show elapsed time instead/as well as ETA; nicer
 
36
# when the rate is unpredictable
 
37
 
 
38
 
 
39
import sys
 
40
import time
 
41
import os
 
42
 
 
43
import bzrlib.errors as errors
 
44
from bzrlib.trace import mutter
 
45
 
 
46
 
 
47
def _supports_progress(f):
 
48
    isatty = getattr(f, 'isatty', None)
 
49
    if isatty is None:
 
50
        return False
 
51
    if not isatty():
 
52
        return False
 
53
    if os.environ.get('TERM') == 'dumb':
 
54
        # e.g. emacs compile window
 
55
        return False
 
56
    return True
 
57
 
 
58
 
 
59
_progress_bar_types = {}
 
60
 
 
61
 
 
62
def ProgressBar(to_file=None, **kwargs):
 
63
    """Abstract factory"""
 
64
    if to_file is None:
 
65
        to_file = sys.stderr
 
66
    requested_bar_type = os.environ.get('BZR_PROGRESS_BAR')
 
67
    # An value of '' or not set reverts to standard processing
 
68
    if requested_bar_type in (None, ''):
 
69
        if _supports_progress(to_file):
 
70
            return TTYProgressBar(to_file=to_file, **kwargs)
 
71
        else:
 
72
            return DotsProgressBar(to_file=to_file, **kwargs)
 
73
    else:
 
74
        # Minor sanitation to prevent spurious errors
 
75
        requested_bar_type = requested_bar_type.lower().strip()
 
76
        # TODO: jam 20060710 Arguably we shouldn't raise an exception
 
77
        #       but should instead just disable progress bars if we
 
78
        #       don't recognize the type
 
79
        if requested_bar_type not in _progress_bar_types:
 
80
            raise errors.InvalidProgressBarType(requested_bar_type,
 
81
                                                _progress_bar_types.keys())
 
82
        return _progress_bar_types[requested_bar_type](to_file=to_file, **kwargs)
 
83
 
 
84
 
 
85
class ProgressBarStack(object):
 
86
    """A stack of progress bars."""
 
87
 
 
88
    def __init__(self,
 
89
                 to_file=None,
 
90
                 show_pct=False,
 
91
                 show_spinner=True,
 
92
                 show_eta=False,
 
93
                 show_bar=True,
 
94
                 show_count=True,
 
95
                 to_messages_file=None,
 
96
                 klass=None):
 
97
        """Setup the stack with the parameters the progress bars should have."""
 
98
        if to_file is None:
 
99
            to_file = sys.stderr
 
100
        if to_messages_file is None:
 
101
            to_messages_file = sys.stdout
 
102
        self._to_file = to_file
 
103
        self._show_pct = show_pct
 
104
        self._show_spinner = show_spinner
 
105
        self._show_eta = show_eta
 
106
        self._show_bar = show_bar
 
107
        self._show_count = show_count
 
108
        self._to_messages_file = to_messages_file
 
109
        self._stack = []
 
110
        self._klass = klass or ProgressBar
 
111
 
 
112
    def top(self):
 
113
        if len(self._stack) != 0:
 
114
            return self._stack[-1]
 
115
        else:
 
116
            return None
 
117
 
 
118
    def bottom(self):
 
119
        if len(self._stack) != 0:
 
120
            return self._stack[0]
 
121
        else:
 
122
            return None
 
123
 
 
124
    def get_nested(self):
 
125
        """Return a nested progress bar."""
 
126
        if len(self._stack) == 0:
 
127
            func = self._klass
 
128
        else:
 
129
            func = self.top().child_progress
 
130
        new_bar = func(to_file=self._to_file,
 
131
                       show_pct=self._show_pct,
 
132
                       show_spinner=self._show_spinner,
 
133
                       show_eta=self._show_eta,
 
134
                       show_bar=self._show_bar,
 
135
                       show_count=self._show_count,
 
136
                       to_messages_file=self._to_messages_file,
 
137
                       _stack=self)
 
138
        self._stack.append(new_bar)
 
139
        return new_bar
 
140
 
 
141
    def return_pb(self, bar):
 
142
        """Return bar after its been used."""
 
143
        if bar is not self._stack[-1]:
 
144
            raise errors.MissingProgressBarFinish()
 
145
        self._stack.pop()
 
146
 
 
147
 
 
148
class _BaseProgressBar(object):
 
149
 
 
150
    def __init__(self,
 
151
                 to_file=None,
 
152
                 show_pct=False,
 
153
                 show_spinner=False,
 
154
                 show_eta=False,
 
155
                 show_bar=True,
 
156
                 show_count=True,
 
157
                 to_messages_file=None,
 
158
                 _stack=None):
 
159
        object.__init__(self)
 
160
        if to_file is None:
 
161
            to_file = sys.stderr
 
162
        if to_messages_file is None:
 
163
            to_messages_file = sys.stdout
 
164
        self.to_file = to_file
 
165
        self.to_messages_file = to_messages_file
 
166
        self.last_msg = None
 
167
        self.last_cnt = None
 
168
        self.last_total = None
 
169
        self.show_pct = show_pct
 
170
        self.show_spinner = show_spinner
 
171
        self.show_eta = show_eta
 
172
        self.show_bar = show_bar
 
173
        self.show_count = show_count
 
174
        self._stack = _stack
 
175
        # seed throttler
 
176
        self.MIN_PAUSE = 0.1 # seconds
 
177
        now = time.clock()
 
178
        # starting now
 
179
        self.start_time = now
 
180
        # next update should not throttle
 
181
        self.last_update = now - self.MIN_PAUSE - 1
 
182
 
 
183
    def finished(self):
 
184
        """Return this bar to its progress stack."""
 
185
        self.clear()
 
186
        assert self._stack is not None
 
187
        self._stack.return_pb(self)
 
188
 
 
189
    def note(self, fmt_string, *args, **kwargs):
 
190
        """Record a note without disrupting the progress bar."""
 
191
        self.clear()
 
192
        self.to_messages_file.write(fmt_string % args)
 
193
        self.to_messages_file.write('\n')
 
194
 
 
195
    def child_progress(self, **kwargs):
 
196
        return ChildProgress(**kwargs)
 
197
 
 
198
 
 
199
class DummyProgress(_BaseProgressBar):
 
200
    """Progress-bar standin that does nothing.
 
201
 
 
202
    This can be used as the default argument for methods that
 
203
    take an optional progress indicator."""
 
204
    def tick(self):
 
205
        pass
 
206
 
 
207
    def update(self, msg=None, current=None, total=None):
 
208
        pass
 
209
 
 
210
    def child_update(self, message, current, total):
 
211
        pass
 
212
 
 
213
    def clear(self):
 
214
        pass
 
215
        
 
216
    def note(self, fmt_string, *args, **kwargs):
 
217
        """See _BaseProgressBar.note()."""
 
218
 
 
219
    def child_progress(self, **kwargs):
 
220
        return DummyProgress(**kwargs)
 
221
 
 
222
 
 
223
_progress_bar_types['dummy'] = DummyProgress
 
224
_progress_bar_types['none'] = DummyProgress
 
225
 
 
226
 
 
227
class DotsProgressBar(_BaseProgressBar):
 
228
 
 
229
    def __init__(self, **kwargs):
 
230
        _BaseProgressBar.__init__(self, **kwargs)
 
231
        self.last_msg = None
 
232
        self.need_nl = False
 
233
        
 
234
    def tick(self):
 
235
        self.update()
 
236
        
 
237
    def update(self, msg=None, current_cnt=None, total_cnt=None):
 
238
        if msg and msg != self.last_msg:
 
239
            if self.need_nl:
 
240
                self.to_file.write('\n')
 
241
            self.to_file.write(msg + ': ')
 
242
            self.last_msg = msg
 
243
        self.need_nl = True
 
244
        self.to_file.write('.')
 
245
        
 
246
    def clear(self):
 
247
        if self.need_nl:
 
248
            self.to_file.write('\n')
 
249
        self.need_nl = False
 
250
        
 
251
    def child_update(self, message, current, total):
 
252
        self.tick()
 
253
 
 
254
 
 
255
_progress_bar_types['dots'] = DotsProgressBar
 
256
 
 
257
    
 
258
class TTYProgressBar(_BaseProgressBar):
 
259
    """Progress bar display object.
 
260
 
 
261
    Several options are available to control the display.  These can
 
262
    be passed as parameters to the constructor or assigned at any time:
 
263
 
 
264
    show_pct
 
265
        Show percentage complete.
 
266
    show_spinner
 
267
        Show rotating baton.  This ticks over on every update even
 
268
        if the values don't change.
 
269
    show_eta
 
270
        Show predicted time-to-completion.
 
271
    show_bar
 
272
        Show bar graph.
 
273
    show_count
 
274
        Show numerical counts.
 
275
 
 
276
    The output file should be in line-buffered or unbuffered mode.
 
277
    """
 
278
    SPIN_CHARS = r'/-\|'
 
279
 
 
280
 
 
281
    def __init__(self, **kwargs):
 
282
        from bzrlib.osutils import terminal_width
 
283
        _BaseProgressBar.__init__(self, **kwargs)
 
284
        self.spin_pos = 0
 
285
        self.width = terminal_width()
 
286
        self.last_updates = []
 
287
        self._max_last_updates = 10
 
288
        self.child_fraction = 0
 
289
        self._have_output = False
 
290
    
 
291
 
 
292
    def throttle(self, old_msg):
 
293
        """Return True if the bar was updated too recently"""
 
294
        # time.time consistently takes 40/4000 ms = 0.01 ms.
 
295
        # but every single update to the pb invokes it.
 
296
        # so we use time.clock which takes 20/4000 ms = 0.005ms
 
297
        # on the downside, time.clock() appears to have approximately
 
298
        # 10ms granularity, so we treat a zero-time change as 'throttled.'
 
299
        now = time.clock()
 
300
        if self.start_time is not None and (now - self.start_time) < 1:
 
301
            return True
 
302
        if old_msg != self.last_msg:
 
303
            return False
 
304
        interval = now - self.last_update
 
305
        # if interval > 0
 
306
        if interval < self.MIN_PAUSE:
 
307
            return True
 
308
 
 
309
        self.last_updates.append(now - self.last_update)
 
310
        # Don't let the queue grow without bound
 
311
        self.last_updates = self.last_updates[-self._max_last_updates:]
 
312
        self.last_update = now
 
313
        return False
 
314
        
 
315
    def tick(self):
 
316
        self.update(self.last_msg, self.last_cnt, self.last_total, 
 
317
                    self.child_fraction)
 
318
 
 
319
    def child_update(self, message, current, total):
 
320
        if current is not None and total != 0:
 
321
            child_fraction = float(current) / total
 
322
            if self.last_cnt is None:
 
323
                pass
 
324
            elif self.last_cnt + child_fraction <= self.last_total:
 
325
                self.child_fraction = child_fraction
 
326
            else:
 
327
                mutter('not updating child fraction')
 
328
        if self.last_msg is None:
 
329
            self.last_msg = ''
 
330
        self.tick()
 
331
 
 
332
    def update(self, msg, current_cnt=None, total_cnt=None, 
 
333
               child_fraction=0):
 
334
        """Update and redraw progress bar."""
 
335
        if msg is None:
 
336
            msg = self.last_msg
 
337
 
 
338
        if total_cnt is None:
 
339
            total_cnt = self.last_total
 
340
 
 
341
        if current_cnt < 0:
 
342
            current_cnt = 0
 
343
            
 
344
        if current_cnt > total_cnt:
 
345
            total_cnt = current_cnt
 
346
        
 
347
        ## # optional corner case optimisation 
 
348
        ## # currently does not seem to fire so costs more than saved.
 
349
        ## # trivial optimal case:
 
350
        ## # NB if callers are doing a clear and restore with
 
351
        ## # the saved values, this will prevent that:
 
352
        ## # in that case add a restore method that calls
 
353
        ## # _do_update or some such
 
354
        ## if (self.last_msg == msg and
 
355
        ##     self.last_cnt == current_cnt and
 
356
        ##     self.last_total == total_cnt and
 
357
        ##     self.child_fraction == child_fraction):
 
358
        ##     return
 
359
 
 
360
        old_msg = self.last_msg
 
361
        # save these for the tick() function
 
362
        self.last_msg = msg
 
363
        self.last_cnt = current_cnt
 
364
        self.last_total = total_cnt
 
365
        self.child_fraction = child_fraction
 
366
 
 
367
        # each function call takes 20ms/4000 = 0.005 ms, 
 
368
        # but multiple that by 4000 calls -> starts to cost.
 
369
        # so anything to make this function call faster
 
370
        # will improve base 'diff' time by up to 0.1 seconds.
 
371
        if self.throttle(old_msg):
 
372
            return
 
373
 
 
374
        if self.show_eta and self.start_time and self.last_total:
 
375
            eta = get_eta(self.start_time, self.last_cnt + self.child_fraction, 
 
376
                    self.last_total, last_updates = self.last_updates)
 
377
            eta_str = " " + str_tdelta(eta)
 
378
        else:
 
379
            eta_str = ""
 
380
 
 
381
        if self.show_spinner:
 
382
            spin_str = self.SPIN_CHARS[self.spin_pos % 4] + ' '            
 
383
        else:
 
384
            spin_str = ''
 
385
 
 
386
        # always update this; it's also used for the bar
 
387
        self.spin_pos += 1
 
388
 
 
389
        if self.show_pct and self.last_total and self.last_cnt:
 
390
            pct = 100.0 * ((self.last_cnt + self.child_fraction) / self.last_total)
 
391
            pct_str = ' (%5.1f%%)' % pct
 
392
        else:
 
393
            pct_str = ''
 
394
 
 
395
        if not self.show_count:
 
396
            count_str = ''
 
397
        elif self.last_cnt is None:
 
398
            count_str = ''
 
399
        elif self.last_total is None:
 
400
            count_str = ' %i' % (self.last_cnt)
 
401
        else:
 
402
            # make both fields the same size
 
403
            t = '%i' % (self.last_total)
 
404
            c = '%*i' % (len(t), self.last_cnt)
 
405
            count_str = ' ' + c + '/' + t 
 
406
 
 
407
        if self.show_bar:
 
408
            # progress bar, if present, soaks up all remaining space
 
409
            cols = self.width - 1 - len(self.last_msg) - len(spin_str) - len(pct_str) \
 
410
                   - len(eta_str) - len(count_str) - 3
 
411
 
 
412
            if self.last_total:
 
413
                # number of markers highlighted in bar
 
414
                markers = int(round(float(cols) * 
 
415
                              (self.last_cnt + self.child_fraction) / self.last_total))
 
416
                bar_str = '[' + ('=' * markers).ljust(cols) + '] '
 
417
            elif False:
 
418
                # don't know total, so can't show completion.
 
419
                # so just show an expanded spinning thingy
 
420
                m = self.spin_pos % cols
 
421
                ms = (' ' * m + '*').ljust(cols)
 
422
                
 
423
                bar_str = '[' + ms + '] '
 
424
            else:
 
425
                bar_str = ''
 
426
        else:
 
427
            bar_str = ''
 
428
 
 
429
        m = spin_str + bar_str + self.last_msg + count_str + pct_str + eta_str
 
430
 
 
431
        assert len(m) < self.width
 
432
        self.to_file.write('\r' + m.ljust(self.width - 1))
 
433
        self._have_output = True
 
434
        #self.to_file.flush()
 
435
            
 
436
    def clear(self):        
 
437
        if self._have_output:
 
438
            self.to_file.write('\r%s\r' % (' ' * (self.width - 1)))
 
439
        self._have_output = False
 
440
        #self.to_file.flush()        
 
441
 
 
442
 
 
443
_progress_bar_types['tty'] = TTYProgressBar
 
444
 
 
445
 
 
446
class ChildProgress(_BaseProgressBar):
 
447
    """A progress indicator that pushes its data to the parent"""
 
448
 
 
449
    def __init__(self, _stack, **kwargs):
 
450
        _BaseProgressBar.__init__(self, _stack=_stack, **kwargs)
 
451
        self.parent = _stack.top()
 
452
        self.current = None
 
453
        self.total = None
 
454
        self.child_fraction = 0
 
455
        self.message = None
 
456
 
 
457
    def update(self, msg, current_cnt=None, total_cnt=None):
 
458
        self.current = current_cnt
 
459
        self.total = total_cnt
 
460
        self.message = msg
 
461
        self.child_fraction = 0
 
462
        self.tick()
 
463
 
 
464
    def child_update(self, message, current, total):
 
465
        if current is None or total == 0:
 
466
            self.child_fraction = 0
 
467
        else:
 
468
            self.child_fraction = float(current) / total
 
469
        self.tick()
 
470
 
 
471
    def tick(self):
 
472
        if self.current is None:
 
473
            count = None
 
474
        else:
 
475
            count = self.current+self.child_fraction
 
476
            if count > self.total:
 
477
                if __debug__:
 
478
                    mutter('clamping count of %d to %d' % (count, self.total))
 
479
                count = self.total
 
480
        self.parent.child_update(self.message, count, self.total)
 
481
 
 
482
    def clear(self):
 
483
        pass
 
484
 
 
485
    def note(self, *args, **kwargs):
 
486
        self.parent.note(*args, **kwargs)
 
487
 
 
488
 
 
489
def str_tdelta(delt):
 
490
    if delt is None:
 
491
        return "-:--:--"
 
492
    delt = int(round(delt))
 
493
    return '%d:%02d:%02d' % (delt/3600,
 
494
                             (delt/60) % 60,
 
495
                             delt % 60)
 
496
 
 
497
 
 
498
def get_eta(start_time, current, total, enough_samples=3, last_updates=None, n_recent=10):
 
499
    if start_time is None:
 
500
        return None
 
501
 
 
502
    if not total:
 
503
        return None
 
504
 
 
505
    if current < enough_samples:
 
506
        return None
 
507
 
 
508
    if current > total:
 
509
        return None                     # wtf?
 
510
 
 
511
    elapsed = time.clock() - start_time
 
512
 
 
513
    if elapsed < 2.0:                   # not enough time to estimate
 
514
        return None
 
515
    
 
516
    total_duration = float(elapsed) * float(total) / float(current)
 
517
 
 
518
    assert total_duration >= elapsed
 
519
 
 
520
    if last_updates and len(last_updates) >= n_recent:
 
521
        avg = sum(last_updates) / float(len(last_updates))
 
522
        time_left = avg * (total - current)
 
523
 
 
524
        old_time_left = total_duration - elapsed
 
525
 
 
526
        # We could return the average, or some other value here
 
527
        return (time_left + old_time_left) / 2
 
528
 
 
529
    return total_duration - elapsed
 
530
 
 
531
 
 
532
class ProgressPhase(object):
 
533
    """Update progress object with the current phase"""
 
534
    def __init__(self, message, total, pb):
 
535
        object.__init__(self)
 
536
        self.pb = pb
 
537
        self.message = message
 
538
        self.total = total
 
539
        self.cur_phase = None
 
540
 
 
541
    def next_phase(self):
 
542
        if self.cur_phase is None:
 
543
            self.cur_phase = 0
 
544
        else:
 
545
            self.cur_phase += 1
 
546
        assert self.cur_phase < self.total 
 
547
        self.pb.update(self.message, self.cur_phase, self.total)
 
548
 
 
549
 
 
550
def run_tests():
 
551
    import doctest
 
552
    result = doctest.testmod()
 
553
    if result[1] > 0:
 
554
        if result[0] == 0:
 
555
            print "All tests passed"
 
556
    else:
 
557
        print "No tests to run"
 
558
 
 
559
 
 
560
def demo():
 
561
    sleep = time.sleep
 
562
    
 
563
    print 'dumb-terminal test:'
 
564
    pb = DotsProgressBar()
 
565
    for i in range(100):
 
566
        pb.update('Leoparden', i, 99)
 
567
        sleep(0.1)
 
568
    sleep(1.5)
 
569
    pb.clear()
 
570
    sleep(1.5)
 
571
    
 
572
    print 'smart-terminal test:'
 
573
    pb = ProgressBar(show_pct=True, show_bar=True, show_spinner=False)
 
574
    for i in range(100):
 
575
        pb.update('Elephanten', i, 99)
 
576
        sleep(0.1)
 
577
    sleep(2)
 
578
    pb.clear()
 
579
    sleep(1)
 
580
 
 
581
    print 'done!'
 
582
 
 
583
if __name__ == "__main__":
 
584
    demo()