~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/progress.py

  • Committer: mbp at sourcefrog
  • Date: 2005-04-11 02:53:57 UTC
  • Revision ID: mbp@sourcefrog.net-20050411025357-af577721308648ae
- remove profiler temporary file when done

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