pgtoolkit/progresswrapper.py
changeset 34 98c7809af415
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/pgtoolkit/progresswrapper.py	Thu May 10 08:42:21 2012 +0200
@@ -0,0 +1,54 @@
+
+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()
+