From: Nikos Skalkotos Date: Mon, 23 Apr 2012 15:53:49 +0000 (+0300) Subject: Use progress for progress bar & implement register X-Git-Tag: v0.1~103 X-Git-Url: https://code.grnet.gr/git/snf-image-creator/commitdiff_plain/b13959674094c57ea0b5dce1164da4819bbcfe3c Use progress for progress bar & implement register --- diff --git a/image_creator/disk.py b/image_creator/disk.py index 1c9d04e..e0ed6a1 100644 --- a/image_creator/disk.py +++ b/image_creator/disk.py @@ -179,16 +179,17 @@ class DiskDevice(object): def enable(self): """Enable a newly created DiskDevice""" - new_progress = progress("Launching helper VM: ") - self.progressbar = new_progress() - self.progressbar.next() + self.progressbar = progress("Launching helper VM: ", "percent") + self.progressbar.max = 100 + self.progressbar.goto(1) eh = self.g.set_event_callback(self.progress_callback, guestfs.EVENT_PROGRESS) self.g.launch() self.guestfs_enabled = True self.g.delete_event_callback(eh) if self.progressbar is not None: - self.progressbar.send(100) + output("\rLaunching helper VM...\033[K", False) + success("done") self.progressbar = None output('Inspecting Operating System...', False) @@ -217,10 +218,7 @@ class DiskDevice(object): position = array[2] total = array[3] - self.progressbar.send((position * 100) // total) - - if position == total: - self.progressbar = None + self.progressbar.goto((position * 100) // total) def mount(self): """Mount all disk partitions in a correct order.""" @@ -318,8 +316,8 @@ class DiskDevice(object): blocksize = 2 ** 22 # 4MB size = self.size() progress_size = (size + 2 ** 20 - 1) // 2 ** 20 # in MB - new_progress = progress("Dumping image file: ") - progressbar = new_progress(progress_size) + progressbar = progress("Dumping image file: ", 'mb') + progressbar.max = progress_size source = open(self.device, "r") try: dest = open(outfile, "w") @@ -333,13 +331,13 @@ class DiskDevice(object): length) offset += sent left -= sent - for i in range((length + 2 ** 20 - 1) // 2 ** 20): - progressbar.next() + progressbar.goto((size - left) // 2 ** 20) finally: dest.close() finally: source.close() - - success('Image file %s was successfully created' % outfile) + + output("\rDumping image file...\033[K", False) + success('image file %s was successfully created' % outfile) # vim: set sta sts=4 shiftwidth=4 sw=4 et ai : diff --git a/image_creator/kamaki_wrapper.py b/image_creator/kamaki_wrapper.py index 2c0db38..f325be0 100644 --- a/image_creator/kamaki_wrapper.py +++ b/image_creator/kamaki_wrapper.py @@ -37,12 +37,30 @@ from kamaki.config import Config from kamaki.clients import ClientError from kamaki.clients.image import ImageClient from kamaki.clients.pithos import PithosClient +from progress.bar import Bar -from image_creator.util import FatalError, progress +from image_creator.util import FatalError, output, success CONTAINER = "images" +def progress(message): + + MSG_LENGTH = 30 + + def progress_gen(n): + msg = "%s:" % message + + progressbar = Bar(msg.ljust(MSG_LENGTH)) + progressbar.max = n + for _ in range(n): + yield + progressbar.next() + output("\r%s...\033[K" % message, False) + success("done") + yield + return progress_gen + class Kamaki: def __init__(self, account, token): self.account = account @@ -60,30 +78,32 @@ class Kamaki: self.uploaded_object = None - def upload(self, filename, size=None, remote_path=None): + def upload(self, file_obj, size=None, remote_path=None, hp=None, up=None): if remote_path is None: remote_path = basename(filename) - with open(filename) as f: - try: - self.pithos_client.create_container(self.container) - except ClientError as e: - if e.status != 202: # Ignore container already exists errors - raise FatalError("Pithos client: %d %s" % \ - (e.status, e.message)) - try: - hash_progress = progress("(1/2) Calculating block hashes:") - upload_progress = progress("(2/2) Uploading missing blocks:") - self.pithos_client.create_object(remote_path, f, size, - hash_progress, upload_progress) - self.uploaded_object = "pithos://%s/%s/%s" % \ - (self.account, self.container, remote_path) - except ClientError as e: + try: + self.pithos_client.create_container(self.container) + except ClientError as e: + if e.status != 202: # Ignore container already exists errors raise FatalError("Pithos client: %d %s" % \ - (e.status, e.message)) + (e.status, e.message)) + try: + hash_cb = progress(hp) if hp is not None else None + upload_cb = progress(up) if up is not None else None + self.pithos_client.create_object(remote_path, file_obj, size, + hash_cb, upload_cb) + return "pithos://%s/%s/%s" % \ + (self.account, self.container, remote_path) + except ClientError as e: + raise FatalError("Pithos client: %d %s" % (e.status, e.message)) - def register(self, metadata): - pass + def register(self, name, location, metadata): + params = {'is_public':'true', 'disk_format':'diskdump'} + try: + self.image_client.register(name, location, params, metadata) + except ClientError as e: + raise FatalError("Image client: %d %s" % (e.status, e.message)) # vim: set sta sts=4 shiftwidth=4 sw=4 et ai : diff --git a/image_creator/main.py b/image_creator/main.py index af236d6..3207c87 100755 --- a/image_creator/main.py +++ b/image_creator/main.py @@ -43,6 +43,7 @@ from image_creator.kamaki_wrapper import Kamaki import sys import os import optparse +import StringIO dd = get_command('dd') @@ -194,34 +195,61 @@ def image_creator(): size = options.shrink and dev.shrink() or dev.size() metadata['SIZE'] = str(size // 2 ** 20) - #Calculating MD5sum - output("Calculating md5sum...", False) checksum = md5(snapshot, size) - success(checksum) - output() - if options.outfile is not None: - f = open('%s.%s' % (options.outfile, 'meta'), 'w') - try: - for key in metadata.keys(): - f.write("%s=%s\n" % (key, metadata[key])) - finally: - f.close() + metastring = "\n".join( + ['%s=%s' % (key, value) for (key, value) in metadata.items()]) + if options.outfile is not None: dev.dump(options.outfile) + output('Dumping metadata file...', False) + with open('%s.%s' % (options.outfile, 'meta'), 'w') as f: + f.write(metastring) + success('done') + + output('Dumping md5sum file...', False) + with open('%s.%s' % (options.outfile, 'md5sum'), 'w') as f: + f.write('%s %s'% (checksum, os.path.basename(options.outfile))) + success('done') + # Destroy the device. We only need the snapshot from now on disk.destroy_device(dev) + output() + + uploaded_obj = "" if options.upload: output("Uploading image to pithos:") kamaki = Kamaki(options.account, options.token) - kamaki.upload(snapshot, size, options.upload) + with open(snapshot) as f: + uploaded_obj = kamaki.upload(f, size, options.upload, + "(1/4) Calculating block hashes", + "(2/4) Uploading missing blocks") + + output("(3/4) Uploading metadata file...", False) + kamaki.upload(StringIO.StringIO(metastring), size=len(metastring), + remote_path="%s.%s" % (options.upload, 'meta')) + success('done') + output("(4/4) Uploading md5sum file...", False) + md5sumstr = '%s %s' % (checksum, os.path.basename(options.upload)) + kamaki.upload(StringIO.StringIO(md5sumstr), size=len(md5sumstr), + remote_path="%s.%s" % (options.upload, 'md5sum')) + success('done') + output() + + if options.register: + output('Registing image to ~okeanos...') + kamaki.register(options.register, uploaded_obj, metadata) + output('done') + output() finally: output('cleaning up...') disk.cleanup() + success("snf-image-creator exited without errors") + return 0 diff --git a/image_creator/util.py b/image_creator/util.py index 718b819..3cacba0 100644 --- a/image_creator/util.py +++ b/image_creator/util.py @@ -34,7 +34,8 @@ import sys import pbs import hashlib -from clint.textui import colored, progress as uiprogress +from clint.textui import colored +from progress.bar import Bar class FatalError(Exception): @@ -85,26 +86,25 @@ def output(msg="", new_line=True): sys.stdout.flush() -def progress(message=''): +def progress(message='', bar_type="default"): - PROGRESS_LENGTH = 32 - MESSAGE_LENGTH = 32 + MESSAGE_LENGTH = 30 + + suffix={'default':'%(index)d/%(max)d', + 'percent':'%(percent)d%%', + 'b':'%(index)d/%(max)d B', + 'kb':'%(index)d/%(max)d KB', + 'mb':'%(index)d/%(max)d MB'} - def progress_generator(n=100): - position = 0 - msg = message.ljust(MESSAGE_LENGTH) - for i in uiprogress.bar(range(n), msg, PROGRESS_LENGTH, silent): - if i < position: - continue - position = yield - yield # suppress the StopIteration exception - return progress_generator + return Bar(message=message.ljust(MESSAGE_LENGTH), fill='#', \ + suffix=suffix[bar_type]) +def md5(filename, size): -def md5(filename, size, progress=None): - - BLOCKSIZE = 2 ^ 22 # 4MB + BLOCKSIZE = 2 ** 22 # 4MB + progressbar = progress("Calculating md5sum:", 'mb') + progressbar.max = (size // (2 ** 20)) md5 = hashlib.md5() with open(filename, "r") as src: left = size @@ -113,7 +113,12 @@ def md5(filename, size, progress=None): data = src.read(length) md5.update(data) left -= length + progressbar.goto((size - left) // (2 ** 20)) + + checksum = md5.hexdigest() + output("\rCalculating md5sum...\033[K", False) + success(checksum) - return md5.hexdigest() + return checksum # vim: set sta sts=4 shiftwidth=4 sw=4 et ai : diff --git a/setup.py b/setup.py index ea41ada..29f0cd5 100755 --- a/setup.py +++ b/setup.py @@ -47,7 +47,7 @@ setup( license='BSD', packages=['image_creator', 'image_creator.os_type'], include_package_data=True, - install_requires=['pbs', 'clint', 'pysendfile'], + install_requires=['pbs', 'clint', 'progress','pysendfile'], entry_points={ 'console_scripts': ['snf-image-creator = image_creator.main:main'] }