Progress

progress.py

import sys
from datetime import datetime, timedelta

try:
    from IPython.display import clear_output
except ImportError:
    clear_output = None


class Progress(object):
    """A context manager for keeping track of the progress of a set
    of tasks. It should be used with the `with` statement, for example:
    
        import time

        # create the progress object with 10 tasks
        p = Progress(10)
        
        for i in range(10):
            # do the task within the context of the progress object
            with p:
                time.sleep(1)
    
    """
    
    def __init__(self, total_tasks):
        # total number of tasks that need to be completed
        self.total_tasks = total_tasks
        # the number of tasks that have been completed so far
        self.task_count = 0
        # the total time elapsed
        self.total_time = timedelta()
        # a running average of time per task
        self.average_time = None
        # estimated time remaining, based on average time per task
        self.time_remaining = None
        # percent of tasks remaining
        self.percent_remaining = None
        
        # a format string for printing out the progress
        self.fmt = "{task_count:d}/{total_tasks:d} tasks completed in {total_time:s}; {percent_remaining:.1%} remaining ({time_remaining:s})"
    
        # the start time of the current task 
        self._time = None
        self._end_time = None
    
    def __enter__(self):
        """Record the time before the task is executed."""
        self._time = datetime.now()

    def __exit__(self, type, value, tb):
        """Compute and print progress after the task is executed."""
        if type is not None:
            return

        # update progress
        self.total_time += datetime.now() - self._time
        self.task_count += 1
        
        # recompute averages and estimates
        self.average_time = self.total_time / self.task_count
        self.time_remaining = (self.total_tasks - self.task_count) * self.average_time
        self.percent_remaining = (self.total_tasks - self.task_count) / self.total_tasks

        # clear existing output and print progress
        if clear_output:
            clear_output(wait=True)
        print(self)
        sys.stdout.flush()
    
    def __str__(self):
        """String representation of the progress object, which is the progress itself."""
        return self.fmt.format(
            task_count=self.task_count,
            total_tasks=self.total_tasks,
            total_time=str(self.total_time),
            percent_remaining=self.percent_remaining,
            time_remaining=str(self.time_remaining))