progress.py 8.32 KB
Newer Older
1 2 3 4 5
# -*- coding: utf-8 -*-
#
# diffoscope: in-depth comparison of files, archives, and directories
#
# Copyright © 2016 Chris Lamb <lamby@debian.org>
6
# Copyright © 2018 Paul Wise <pabs@debian.org>
7 8 9 10 11 12 13 14 15 16 17 18
#
# 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
19
# along with diffoscope.  If not, see <https://www.gnu.org/licenses/>.
20

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

27
from .logging import line_eraser
28 29


30 31
logger = logging.getLogger(__name__)

32

33
class ProgressLoggingHandler(logging.StreamHandler):
34 35
    def __init__(self, progressbar):
        self.progressbar = progressbar
36
        super().__init__()
37 38

    def emit(self, record):
39 40 41 42 43
        try:
            super().emit(record)
            if not self.progressbar.bar.finished:
                self.progressbar.bar.update()
        except Exception:
44 45
            # Wrap otherwise tests fail due to test_progress.py call main()
            # several times. This mirrors the super() implementation.
46
            self.handleError(record)
47

48

49 50 51 52 53 54 55
class ProgressManager(object):
    _singleton = {}

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

        if not self._singleton:
56 57 58
            self.reset()

    def reset(self):
59
        self.stack = []
60
        self.observers = []
61 62

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

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

            return False

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

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

88
        return log_handler
89

90 91 92 93 94 95 96 97 98 99
    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)

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

104 105 106 107 108 109 110 111
    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
112 113

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

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

120

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

    def __enter__(self):
133
        ProgressManager().push(self)
134 135 136
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
137 138 139 140 141 142 143 144 145 146
        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:
147 148
                # Something is in-progress, the calculation is slightly more
                # complex
149 150
                cur_child_done, cur_child_total = cur_child_estimate or (0, 0)
                own_done += self.current_steps
151 152 153 154 155 156 157 158
                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
159
                return all_done, int(float(self.total) / own_done
Mattia Rizzolo's avatar
Mattia Rizzolo committed
160
                                     * expected_all_done)
161 162 163 164
        else:
            # nothing in progress
            assert not cur_child_estimate

165 166 167 168
        if not own_done:
            assert not children_done
            return 0, self.total

169 170 171 172
        # 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
173
        return self.current_steps is not None
174 175 176 177

    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
178 179
            self.current_steps = None
            self.current_child_steps_done = None
180 181 182
            ProgressManager().update(msg)

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

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

191

192 193 194 195
class ProgressBar(object):
    def __init__(self):
        import progressbar

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

202 203
        self.msg = ""

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

                if len(msg) <= width:
210
                    return msg.rjust(width)
211 212 213 214

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

215
        class OurProgressBar(progressbar.ProgressBar):
216 217
            def __init__(self, *args, **kwargs):
                # Remove after https://github.com/niltonvolpato/python-progressbar/pull/57 is fixed.
218
                kwargs.setdefault('fd', sys.stderr)
219
                super().__init__(*args, **kwargs)
220
                # Terminal handling after parent init since that sets self.fd
221
                self.erase_to_eol = line_eraser(self.fd)
222

223 224 225
            def _need_update(self):
                return True

226 227 228
            def erase_line(self):
                if self.erase_to_eol:
                    self.fd.buffer.write(self.erase_to_eol)
229
                    self.fd.flush()
230 231 232 233 234 235 236 237 238

            def finish(self):
                self.finished = True
                self.update(self.maxval)
                # Clear the progress bar after completion
                self.erase_line()
                if self.signal_set:
                    signal.signal(signal.SIGWINCH, signal.SIG_DFL)

239
        self.bar = OurProgressBar(widgets=(
240
            ' ',
241 242 243 244
            progressbar.Bar(),
            '  ',
            progressbar.Percentage(),
            '  ',
245
            Message(),
246
            '  ',
247
            progressbar.ETA(),
248
            ' ',
249 250 251
        ))
        self.bar.start()

252 253 254
    def notify(self, current, total, msg):
        self.msg = msg

255 256 257 258 259 260 261
        self.bar.maxval = total
        self.bar.currval = current
        self.bar.update()

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

262

263
class StatusFD(object):
264 265
    def __init__(self, fileobj):
        self.fileobj = fileobj
266

267
    def notify(self, current, total, msg):
268
        print(json.dumps({
269
            'msg': msg,
270 271 272
            'total': total,
            'current': current,
        }), file=self.fileobj)
273 274 275

    def finish(self):
        pass