consoleLog / src / consoleLog / progressBar.py

# -*- coding: utf-8 -*-
# 
#  progressBar.py
#  consoleLog
#  
#  Created by Lars Yencken on 2008-06-12.
#  Copyright 2008-06-12 Lars Yencken. All rights reserved.
#

"""
A simple console progress bar, with support for iterator wrapping.
"""

#----------------------------------------------------------------------------#

import sys
import time

import shellColor
import slot

#----------------------------------------------------------------------------#

_currentBar = None

class IllegalNestingError(Exception):
    pass

def withProgress(seq, modValue=1, size=None, force=False):
    """
    A wrapper for sequence which adds a progress bar.
    """
    global _currentBar
    if _currentBar:
        raise IllegalNestingError("can't nest progress bars")

    progress = ProgressBar()
    if size is None:
        if hasattr(seq, '__len__'):
            progress.start(len(seq))
        elif force:
            seq = list(seq)
            progress.start(len(seq))
        else:
            raise ValueError, "seq has no __len__ method"
    else:
        progress.start(size)

    _currentBar = progress

    count = 0
    for item in seq:
        yield item
        count += 1
        if count % modValue == 0:
            progress.update(count)

    progress.finish()

    _currentBar = None
    return

#----------------------------------------------------------------------------#

class ProgressIterator(object):
    """
    A context manager which provides an iterator, as well as a progressBar which is
    updated according to the progress through the given sequence.
    """
    def __init__(self, logger, seq, msg):
        self.logger = logger
        self.seq = seq
        self.start_message = msg if msg is not None else ""

    def __enter__(self):
        self.logger.log(self.start_message, newLine=False)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

    def __iter__(self):
        for item in withProgress(self.seq):
          yield item

#----------------------------------------------------------------------------#

def tick():
    """
    Indicate activity in the current progress bar.
    """
    global _currentBar

    if _currentBar is not None:
        _currentBar.tick()
    return

#----------------------------------------------------------------------------#

class ProgressBar:
    "A progress bar which gets printed to stdout."
    def __init__(self, stringSize=20):
        "Creates a new instance, setting the size as needed."
        self._slot = slot.Slot()
        self._stringSize = stringSize
        self._count = 0
        self._totalCount = None
        self._lastRotation = 0
        self._startTime = None

        self._rotation = ['/', '-', '\\', '|']
        self._numRotations = len(self._rotation)

        return

    def reset(self):
        "Resets the progress bar to initial conditions."
        self._count = 0
        self._totalCount = None
        self._lastLineSize = None
        self._lastRotation = 0
        self._startTime = None

        return
    
    def start(self, totalCount):
        """
        Starts the progress bar. This will be the first output from the
        bar. At this stage it needs the total count so that it can
        calculate percentages.

        @param totalCount: The count which represents 'finished'.
        """
        assert self._count == 0, "Progress bar already started, call reset()"
        assert totalCount > 0, "Progress bar needs a non-zero total count"
        self._totalCount = totalCount
        progressLine = '[' + shellColor.color('/', 'red') + \
                (self._stringSize-1)*' ' + ']   0%'
        self._slot.update(progressLine)
        self._startTime = time.time()

    def update(self, count):
        """
        Updates the progress bar with the current count. This is useful
        to call even if the count has not increased, since it provides
        visual feedback to the user that the process is active.

        @param count: The current count
        """
        count = int(count)
        self._count = count
        if count < 0 or count > self._totalCount:
            raise Exception, 'Bad count for progress bar'

        n = (count * self._stringSize) / self._totalCount
        percent = (100*count) / self._totalCount
        m = self._stringSize - n - 1 # subtract one for the rotating char

        self._lastRotation = (self._lastRotation + 1) % self._numRotations
        if percent < 100:
            rotChar = self._rotation[self._lastRotation]
        else:
            rotChar = ''

        progressLine = '[' + shellColor.color(n*'-' + rotChar, 'red') + \
                m*' ' + '] ' + str('%3d%%' % percent)
        self._slot.update(progressLine)

    def tick(self):
        "Perform a quick update."
        self.update(self._count)

    def fractional(self, fraction):
        "Set a fractional percentage completion, e.g. 0.3333 -> 33%."
        assert fraction >= 0 and fraction <= 1
        self.update(int(fraction * self._totalCount))
    
    def finish(self):
        "Fixes to 100% complete, and writes the time taken."
        assert self._totalCount > 0, "Progress bar wasn't initialised"
        self.update(self._totalCount)

        timeTaken = int(time.time() - self._startTime)
        (mins, secs) = divmod(timeTaken, 60)
        if not mins:
            timeString = ' (%ds)' % secs
        else:
            (hours, mins) = divmod(mins, 60)
            if not hours:
                timeString = ' (%dm %ds)' % (mins, secs)
            else:
                timeString = ' (%dh %dm %ds)' % (hours, mins, secs)

        print >> sys.stdout, timeString

#----------------------------------------------------------------------------#

if __name__ == '__main__':
    for i in withProgress(range(100)):
        time.sleep(0.1)
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.