163237
#!/usr/bin/python3
163237
# -*- coding: utf-8 -*-
163237
#
163237
# A simple text-based progress bar, compatible with the basic API of:
163237
# https://github.com/WoLpH/python-progressbar
163237
#
163237
# Copyright (C) 2021  Red Hat, Inc.
163237
# 
163237
# This program is free software; you can redistribute it and/or
163237
# modify it under the terms of the GNU General Public License
163237
# as published by the Free Software Foundation; either version 2
163237
# of the License, or (at your option) any later version.
163237
# 
163237
# This program is distributed in the hope that it will be useful,
163237
# but WITHOUT ANY WARRANTY; without even the implied warranty of
163237
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
163237
# GNU General Public License for more details.
163237
# 
163237
# You should have received a copy of the GNU General Public License
163237
# along with this program; if not, write to the Free Software
163237
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301,
163237
# USA.
163237
163237
163237
import shutil
163237
import sys
163237
import time
163237
163237
163237
class ProgressBar:
163237
    FORMAT = '{value:>10} / {max_value:<10} [{bars}]'
163237
    BARS = '= '
163237
    SPINLEN = 5
163237
163237
    def __init__(self, stream=sys.stderr, max_width=80, fps=10):
163237
        self._stream = stream
163237
        self._max_width = max_width
163237
        self._min_delay = 1 / fps
163237
163237
    @staticmethod
163237
    def _format_value(value):
163237
        raise NotImplementedError()
163237
163237
    def start(self, max_value):
163237
        self._value = 0
163237
        self._max_value = max_value or 0
163237
        self._status = dict()
163237
        self._spinner = 0
163237
        self._timestamp = 0
163237
        self.update(0)
163237
163237
    def update(self, value):
163237
        self._value = value
163237
        if value > self._max_value:
163237
            self._max_value = 0
163237
163237
        ts = time.time()
163237
        if (ts - self._timestamp) < self._min_delay:
163237
            return
163237
        self._timestamp = ts
163237
163237
        status = {'value': self._format_value(value),
163237
                  'max_value': self._format_value(self._max_value) \
163237
                               if self._max_value else '???',
163237
                  'bars': ''}
163237
163237
        termw = min(shutil.get_terminal_size()[0], self._max_width)
163237
        nbars = max(termw - len(self.FORMAT.format(**status)), 0)
163237
        nfill = nskip = 0
163237
163237
        if self._max_value:
163237
            nfill = round(nbars * value / self._max_value)
163237
        elif nbars > self.SPINLEN:
163237
            nfill = self.SPINLEN
163237
            nskip = self._spinner % (nbars - self.SPINLEN)
163237
            self._spinner = nskip + 1
163237
163237
        status['bars'] = self.BARS[1] * nskip + \
163237
                         self.BARS[0] * nfill + \
163237
                         self.BARS[1] * (nbars - nfill - nskip)
163237
163237
        if status == self._status:
163237
            return
163237
        self._status = status
163237
163237
        self._stream.write('\r')
163237
        self._stream.write(self.FORMAT.format(**self._status))
163237
        self._stream.flush()
163237
163237
    def finish(self):
163237
        self._max_value = self._value
163237
        self._timestamp = 0  # Force an update
163237
        self.update(self._value)
163237
163237
        self._stream.write('\n')
163237
        self._stream.flush()
163237
163237
163237
class DataTransferBar(ProgressBar):
163237
    @staticmethod
163237
    def _format_value(value):
163237
        symbols = ' KMGTPEZY'
163237
        depth = 0
163237
        max_depth = len(symbols) - 1
163237
        unit = 1024.0
163237
163237
        # 1023.95 should be formatted as 1.0 (not 1024.0)
163237
        # More info: https://stackoverflow.com/a/63839503
163237
        thres = unit - 0.05
163237
163237
        while value >= thres and depth < max_depth:
163237
            depth += 1
163237
            value /= unit
163237
        symbol = ' %siB' % symbols[depth] if depth > 0 else ''
163237
163237
        return '%.1f%s' % (value, symbol)
163237
163237
163237
if __name__ == '__main__':
163237
    # Show a dummy bar for debugging purposes
163237
163237
    bar = DataTransferBar()
163237
    size = 50*1024*1024
163237
    chunk = 1024*1234
163237
    recvd = 0
163237
163237
    bar.start(size)
163237
    while recvd < (size - chunk):
163237
        recvd += chunk
163237
        bar.update(recvd)
163237
        time.sleep(0.1)
163237
    bar.update(size)
163237
    bar.finish()