Commit | Line | Data |
---|---|---|
418527ca JD |
1 | import os |
2 | import sys | |
3 | ||
4 | try: | |
5 | from progressbar import ETA, Bar, Percentage, ProgressBar | |
6 | progressbar_available = True | |
7 | except ImportError: | |
8 | progressbar_available = False | |
9 | ||
10 | # approximation for the progress bar | |
11 | BYTES_PER_EVENT = 30 | |
12 | ||
13 | ||
14 | def getFolderSize(folder): | |
15 | total_size = os.path.getsize(folder) | |
16 | for item in os.listdir(folder): | |
17 | itempath = os.path.join(folder, item) | |
18 | if os.path.isfile(itempath): | |
19 | total_size += os.path.getsize(itempath) | |
20 | elif os.path.isdir(itempath): | |
21 | total_size += getFolderSize(itempath) | |
22 | return total_size | |
23 | ||
24 | ||
25 | def progressbar_setup(obj): | |
26 | if hasattr(obj, "_arg_no_progress") and obj._arg_no_progress: | |
27 | obj.pbar = None | |
28 | return | |
29 | ||
30 | if progressbar_available: | |
31 | size = getFolderSize(obj._arg_path) | |
32 | widgets = ['Processing the trace: ', Percentage(), ' ', | |
33 | Bar(marker='#', left='[', right=']'), | |
34 | ' ', ETA(), ' '] # see docs for other options | |
35 | obj.pbar = ProgressBar(widgets=widgets, | |
36 | maxval=size/BYTES_PER_EVENT) | |
37 | obj.pbar.start() | |
38 | else: | |
39 | print("Warning: progressbar module not available, " | |
40 | "using --no-progress.", file=sys.stderr) | |
41 | obj._arg_no_progress = True | |
42 | obj.pbar = None | |
43 | obj.event_count = 0 | |
44 | ||
45 | ||
46 | def progressbar_update(obj): | |
47 | if hasattr(obj, "_arg_no_progress") and \ | |
48 | (obj._arg_no_progress or obj.pbar is None): | |
49 | return | |
50 | try: | |
51 | obj.pbar.update(obj.event_count) | |
52 | except ValueError: | |
53 | pass | |
54 | obj.event_count += 1 | |
55 | ||
56 | ||
57 | def progressbar_finish(obj): | |
58 | if hasattr(obj, "_arg_no_progress") and obj._arg_no_progress: | |
59 | return | |
60 | obj.pbar.finish() |