~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/progress.py

  • Committer: mbp at sourcefrog
  • Date: 2005-03-21 22:29:49 UTC
  • Revision ID: mbp@sourcefrog.net-20050321222949-232c2093a6eadd80
fixup doctest for new module structure

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()