pgtoolkit/progresswrapper.py
author Radek Brich <brich.radek@ifortuna.cz>
Tue, 06 May 2014 18:34:38 +0200
changeset 99 245646538743
parent 34 98c7809af415
permissions -rw-r--r--
Update runquery tool: Add --one-query-per-line parameter.


import sys


class ProgressWrapper:
    
    def __init__(self, f, size=0):
        self.f = f
        self.size = size
        self.readsize = 0
        self.cycles = 0
        self.print_cycles = 200

    def humanize(self, bytes):
        if bytes > 1024**3:
            return '%.1fG' % (bytes / 1024.**3)
        if bytes > 1024**2:
            return '%.1fM' % (bytes / 1024.**2)
        if bytes > 1024:
            return '%.1fk' % (bytes / 1024.)
        return '%d' % bytes

    def write(self, data):
        self.size += len(data)
        if self.cycles == 0:
            print('  read %s      \r' % self.humanize(self.size), end='')
            sys.stdout.flush()
            self.cycles = self.print_cycles * 200
        else:
            self.cycles -= 1
        return self.f.write(data)

    def read(self, size):
        self.readsize += size
        if self.cycles == 0:
            if self.size > 0:
                percent = self.readsize * 100. / self.size
            else:
                percent = 100
            if percent > 100:
                percent = 100
            print('  written %s / %s (%.1f%%)      \r' % (
                self.humanize(self.readsize),
                self.humanize(self.size),
                percent), end='')
            sys.stdout.flush()
            self.cycles = self.print_cycles
        else:
            self.cycles -= 1
        return self.f.read(size)

    def close(self):
        self.f.close()