Commits

Josh VanderLinden  committed 4ab73d9

More refactoring

  • Participants
  • Parent commits 60d84a8

Comments (0)

Files changed (1)

 
 from subprocess import Popen, PIPE, call
 import logging
+import math
 import os
 import re
 import select
 import sys
 import time
 
-PROGRESS_RE = re.compile(r'^(?P<bytes>\d+) bytes \((?P<copied>[^\)]+)\) copied, (?P<time>[^,]+ s), (?P<speed>.*?)$')
+PROGRESS_RE = re.compile(r'^(?P<bytes>\d+) bytes \((?P<copied>[^\)]+)\) copied, (?P<time>[^,]+ s), (?P<speed>.*?)$', re.M)
 
-logging.basicConfig(format='%(message)s',
-                    level=logging.WARN,
-                    stream=sys.stderr)
+logging.basicConfig(format='%(message)s', stream=sys.stderr, level=logging.WARN)
 log = logging.getLogger('pp')
 
 
-def get_status(pid):
-    call(['kill', '-SIGUSR1', str(pid)])
+class DDWrapper(object):
 
-def get_output(pipe):
-    return select.select([pipe.stdout, pipe.stderr], [], [], 1)
+    def __init__(self):
+        self.pipe = None
+        self.filesize = None
 
-def show_progress(pipe, source):
-    line = source.readline()
-    if source == pipe.stdout:
-        sys.stdout.write(line)
-        continue
+    def __call__(self, oflag=None, **kwargs):
+        if oflag is None:
+            log.info('Inserting oflag=direct,sync as an argument for more accurate progress')
+            kwargs['oflag'] = 'direct,sync'
 
-    m = PROGRESS_RE.match(line)
-    if not m:
-        continue
+        # get some information about the source file
+        try:
+            input_file = kwargs['if']
+        except KeyError:
+            log.error('Invalid input file')
+            sys.exit(1)
 
-    log.debug('Matched progress line: %s' % (m.groupdict(),))
-    copied = int(m.group('bytes'))
-    progress = copied / float(filesize) * 100.0
-    p = min(progress, 100)
+        try:
+            log.debug('Attempting to stat: %s' % (input_file,))
+            info = os.stat(input_file)
+        except OSError as exc:
+            log.error(str(exc))
+            sys.exit(2)
 
-    # display the progress bar
-    sys.stdout.write('\r[%s] %s%%, %s copied at %s     ' % (
-        ('#' * int(p / 2)).ljust(50, ' '), int(p),
-        m.group('copied'), m.group('speed')
-    ))
-    sys.stdout.flush()
+        self.filesize = info.st_size
+        log.debug('Detected filesize: %s' % (self.filesize,))
 
-    log.debug('Asking for progress')
-    get_status(pipe.pid)
+        args = ['%s=%s' % pair for pair in kwargs.items()]
+        self.pipe = Popen(['dd'] + args, stdout=PIPE, stderr=PIPE)
+        self.get_status()
 
+        while self.pipe.poll() is None:
+            self.show_progress()
 
-def main():
-    args = sys.argv[1:]
-    if 'oflag=' not in args:
-        log.info('Inserting oflag=direct,sync as an argument for more accurate progress')
-        args.insert(0, 'oflag=direct,sync')
+        self.show_progress(True)
+        sys.stdout.write('\n')
 
-    # get some information about the source file
-    try:
-        input_file = [a.split('=')[1] for a in args if a.startswith('if=')][0]
-    except IndexError:
-        log.error('Invalid input file')
-        sys.exit(1)
-
-    log.debug('Attempting to stat: %s' % (input_file,))
-    try:
-        info = os.stat(input_file)
-    except OSError as exc:
-        log.error(str(exc))
-        sys.exit(2)
+    def get_status(self):
+        call(['kill', '-SIGUSR1', str(self.pipe.pid)], stderr=PIPE)
 
-    filesize = info.st_size
-    log.debug('Detected filesize: %s' % (filesize,))
-
-    p = 0
-    copied = 0
-    pipe = Popen(['dd'] + args, stdout=PIPE, stderr=PIPE)
-    get_status(pipe.pid)
-
-    while pipe.poll() is None:
-        r, w, x = get_output(pipe)
+    def show_progress(self, full=False):
+        r, w, x = select.select([self.pipe.stdout, self.pipe.stderr], [], [], 1)
         for source in r:
-            show_progress(pipe, source)
+            if full:
+                line = source.read()
+            else:
+                line = source.readline()
+
+            m = PROGRESS_RE.search(line)
+            if not m:
+                continue
+
+            log.debug('Matched progress line: %s' % (m.groupdict(),))
+            copied = int(m.group('bytes'))
+            progress = copied / float(self.filesize) * 100.0
+            p = min(progress, 100)
+
+            # display the progress bar
+            sys.stdout.write('\r[%s] %s%%, %s copied at %s     ' % (
+                ('#' * int(math.ceil(p / 2))).ljust(50, ' '), int(p),
+                m.group('copied'), m.group('speed')
+            ))
+            sys.stdout.flush()
+
+            log.debug('Asking for progress')
+            self.get_status()
             time.sleep(0.25)
 
-    r, w, x = get_output(pipe)
-    for source in r:
-        sys.stdout.write(source.read())
 
 if __name__ == '__main__':
     try:
-        main()
+        args = sys.argv[1:]
+        kwargs = dict(arg.split('=', 1) for arg in args)
+
+        ddw = DDWrapper()
+        ddw(**kwargs)
     except (KeyboardInterrupt, SystemExit):
         pass
     except: