progress.py 7.75 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
# -*- coding: utf-8 -*-
#
# diffoscope: in-depth comparison of files, archives, and directories
#
# Copyright © 2016 Chris Lamb <lamby@debian.org>
#
# diffoscope is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# diffoscope is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
18
# along with diffoscope.  If not, see <https://www.gnu.org/licenses/>.
19

20
import os
21
import sys
22
import json
23 24 25 26
import logging

logger = logging.getLogger(__name__)

27

28
class ProgressLoggingHandler(logging.StreamHandler):
29 30
    def __init__(self, progressbar):
        self.progressbar = progressbar
31
        super().__init__()
32 33

    def emit(self, record):
34
        try:
35
            # Delete the current line (i.e. the progress bar)
36 37 38 39 40 41
            self.stream.write("\r\033[K")
            self.flush()
            super().emit(record)
            if not self.progressbar.bar.finished:
                self.progressbar.bar.update()
        except Exception:
42 43
            # Wrap otherwise tests fail due to test_progress.py call main()
            # several times. This mirrors the super() implementation.
44
            self.handleError(record)
45

46

47 48 49 50 51 52 53
class ProgressManager(object):
    _singleton = {}

    def __init__(self):
        self.__dict__ = self._singleton

        if not self._singleton:
54 55 56
            self.reset()

    def reset(self):
57
        self.stack = []
58
        self.observers = []
59 60

    def setup(self, parsed_args):
61 62 63 64 65
        def show_progressbar():
            # Show progress bar if user explicitly asked for it
            if parsed_args.progress:
                return True

66
            # ... otherwise show it if STDOUT is a tty and we are not debugging
67
            if parsed_args.progress is None:
68
                return sys.stdout.isatty() and not parsed_args.debug
69 70 71

            return False

72
        log_handler = None
73
        if show_progressbar():
74
            try:
75 76 77
                bar = ProgressBar()
                self.register(bar)
                log_handler = ProgressLoggingHandler(bar)
78 79 80 81 82
            except ImportError:
                # User asked for bar, so show them the error
                if parsed_args.progress:
                    raise

83
        if parsed_args.status_fd:
84
            self.register(StatusFD(os.fdopen(parsed_args.status_fd, 'w')))
85

86
        return log_handler
87

88 89 90 91 92 93 94 95 96 97
    def push(self, progress):
        assert not self.stack or self.stack[-1].is_active()
        self.stack.append(progress)

    def pop(self, progress):
        x = self.stack.pop()
        assert x is progress
        if self.stack:
            self.stack[-1].child_done(x.total)

98
    def register(self, observer):
99
        logger.debug("Registering %s as a progress observer", observer)
100 101
        self.observers.append(observer)

102 103 104 105 106 107 108 109
    def update(self, msg):
        if self.stack:
            cur_estimates = None
            for progress in reversed(self.stack):
                cur_estimates = progress.estimates(cur_estimates)
            current, total = cur_estimates
        else:
            current, total = 0, 1
110 111

        for x in self.observers:
112
            x.notify(current, total, msg)
113 114 115 116 117

    def finish(self):
        for x in self.observers:
            x.finish()

118

119
class Progress(object):
120 121
    def __init__(self, total=None):
        self.done = []
Ximin Luo's avatar
Ximin Luo committed
122 123
        self.current_steps = None
        self.current_child_steps_done = None
124 125 126 127 128
        if total:
            self.total = total
        else:
            self.total = 1
            self.begin_step(1)
129 130

    def __enter__(self):
131
        ProgressManager().push(self)
132 133 134
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
135 136 137 138 139 140 141 142 143 144
        self.maybe_end()
        ProgressManager().pop(self)

    def estimates(self, cur_child_estimate=None):
        own_done = sum(pair[0] for pair in self.done)
        children_done = sum(pair[1] for pair in self.done)
        all_done = own_done + children_done

        if self.current_steps:
            if self.current_child_steps_done or cur_child_estimate:
145 146
                # Something is in-progress, the calculation is slightly more
                # complex
147 148
                cur_child_done, cur_child_total = cur_child_estimate or (0, 0)
                own_done += self.current_steps
149 150 151 152 153 154 155 156
                all_done += self.current_steps + \
                    self.current_child_steps_done + \
                    cur_child_done
                # Cost of what we expect will have been done, once the current
                # in-progress step plus all of its children, have completed
                expected_all_done = all_done + \
                    (cur_child_total - cur_child_done)
                assert own_done  # non-zero
Mattia Rizzolo's avatar
Mattia Rizzolo committed
157
                return all_done, int(float(self.total) / own_done
Mattia Rizzolo's avatar
Mattia Rizzolo committed
158
                                     * expected_all_done)
159 160 161 162
        else:
            # nothing in progress
            assert not cur_child_estimate

163 164 165 166
        if not own_done:
            assert not children_done
            return 0, self.total

167 168 169 170
        # weigh self.total by (all_done/own_done)
        return all_done, int(float(self.total) / own_done * all_done)

    def is_active(self):
Ximin Luo's avatar
Ximin Luo committed
171
        return self.current_steps is not None
172 173 174 175

    def maybe_end(self, msg=""):
        if self.is_active():
            self.done += [(self.current_steps, self.current_child_steps_done)]
Ximin Luo's avatar
Ximin Luo committed
176 177
            self.current_steps = None
            self.current_child_steps_done = None
178 179 180
            ProgressManager().update(msg)

    def begin_step(self, step, msg=""):
Ximin Luo's avatar
Ximin Luo committed
181
        assert step is not None
182 183
        self.maybe_end(msg)
        self.current_steps = step
Ximin Luo's avatar
Ximin Luo committed
184
        self.current_child_steps_done = 0
185 186 187

    def child_done(self, total):
        self.current_child_steps_done += total
188

189

190 191 192 193
class ProgressBar(object):
    def __init__(self):
        import progressbar

194 195 196 197 198 199
        try:
            from progressbar.widgets import WidgetBase
        except ImportError:
            # Fallback to the older Debian version
            from progressbar import Widget as WidgetBase

200 201
        self.msg = ""

202
        class Message(WidgetBase):
203 204
            def update(self, pbar, _observer=self):
                msg = _observer.msg
205
                width = 25
206 207

                if len(msg) <= width:
208
                    return msg.rjust(width)
209 210 211 212

                # Print the last `width` characters with an ellipsis.
                return '…{}'.format(msg[-width + 1:])

213
        class OurProgressBar(progressbar.ProgressBar):
214 215
            def __init__(self, *args, **kwargs):
                # Remove after https://github.com/niltonvolpato/python-progressbar/pull/57 is fixed.
216
                kwargs.setdefault('fd', sys.stderr)
217 218
                super().__init__(*args, **kwargs)

219 220 221 222
            def _need_update(self):
                return True

        self.bar = OurProgressBar(widgets=(
223
            ' ',
224 225 226 227
            progressbar.Bar(),
            '  ',
            progressbar.Percentage(),
            '  ',
228
            Message(),
229
            '  ',
230
            progressbar.ETA(),
231
            ' ',
232 233 234
        ))
        self.bar.start()

235 236 237
    def notify(self, current, total, msg):
        self.msg = msg

238 239 240 241 242 243 244
        self.bar.maxval = total
        self.bar.currval = current
        self.bar.update()

    def finish(self):
        self.bar.finish()

245

246
class StatusFD(object):
247 248
    def __init__(self, fileobj):
        self.fileobj = fileobj
249

250
    def notify(self, current, total, msg):
251
        print(json.dumps({
252
            'msg': msg,
253 254 255
            'total': total,
            'current': current,
        }), file=self.fileobj)
256 257 258

    def finish(self):
        pass