Revision b1395967

b/image_creator/disk.py
179 179

  
180 180
    def enable(self):
181 181
        """Enable a newly created DiskDevice"""
182
        new_progress = progress("Launching helper VM: ")
183
        self.progressbar = new_progress()
184
        self.progressbar.next()
182
        self.progressbar = progress("Launching helper VM: ", "percent")
183
        self.progressbar.max = 100
184
        self.progressbar.goto(1)
185 185
        eh = self.g.set_event_callback(self.progress_callback,
186 186
                                                    guestfs.EVENT_PROGRESS)
187 187
        self.g.launch()
188 188
        self.guestfs_enabled = True
189 189
        self.g.delete_event_callback(eh)
190 190
        if self.progressbar is not None:
191
            self.progressbar.send(100)
191
            output("\rLaunching helper VM...\033[K", False)
192
            success("done")
192 193
            self.progressbar = None
193 194

  
194 195
        output('Inspecting Operating System...', False)
......
217 218
        position = array[2]
218 219
        total = array[3]
219 220

  
220
        self.progressbar.send((position * 100) // total)
221

  
222
        if position == total:
223
            self.progressbar = None
221
        self.progressbar.goto((position * 100) // total)
224 222

  
225 223
    def mount(self):
226 224
        """Mount all disk partitions in a correct order."""
......
318 316
        blocksize = 2 ** 22  # 4MB
319 317
        size = self.size()
320 318
        progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
321
        new_progress = progress("Dumping image file: ")
322
        progressbar = new_progress(progress_size)
319
        progressbar = progress("Dumping image file: ", 'mb')
320
        progressbar.max = progress_size
323 321
        source = open(self.device, "r")
324 322
        try:
325 323
            dest = open(outfile, "w")
......
333 331
                                                                        length)
334 332
                    offset += sent
335 333
                    left -= sent
336
                    for i in range((length + 2 ** 20 - 1) // 2 ** 20):
337
                        progressbar.next()
334
                    progressbar.goto((size - left) // 2 ** 20)
338 335
            finally:
339 336
                dest.close()
340 337
        finally:
341 338
            source.close()
342

  
343
        success('Image file %s was successfully created' % outfile)
339
        
340
        output("\rDumping image file...\033[K", False)
341
        success('image file %s was successfully created' % outfile)
344 342

  
345 343
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
b/image_creator/kamaki_wrapper.py
37 37
from kamaki.clients import ClientError
38 38
from kamaki.clients.image import ImageClient
39 39
from kamaki.clients.pithos import PithosClient
40
from progress.bar import Bar
40 41

  
41
from image_creator.util import FatalError, progress
42
from image_creator.util import FatalError, output, success
42 43

  
43 44
CONTAINER = "images"
44 45

  
45 46

  
47
def progress(message):
48

  
49
    MSG_LENGTH = 30
50

  
51
    def progress_gen(n):
52
        msg = "%s:" % message
53

  
54
        progressbar = Bar(msg.ljust(MSG_LENGTH))
55
        progressbar.max = n
56
        for _ in range(n):
57
            yield
58
            progressbar.next()
59
        output("\r%s...\033[K" % message, False)
60
        success("done")
61
        yield
62
    return progress_gen
63

  
46 64
class Kamaki:
47 65
    def __init__(self, account, token):
48 66
        self.account = account
......
60 78

  
61 79
        self.uploaded_object = None
62 80

  
63
    def upload(self, filename, size=None, remote_path=None):
81
    def upload(self, file_obj, size=None, remote_path=None, hp=None, up=None):
64 82

  
65 83
        if remote_path is None:
66 84
            remote_path = basename(filename)
67 85

  
68
        with open(filename) as f:
69
            try:
70
                self.pithos_client.create_container(self.container)
71
            except ClientError as e:
72
                if e.status != 202:  # Ignore container already exists errors
73
                    raise FatalError("Pithos client: %d %s" % \
74
                                                        (e.status, e.message))
75
            try:
76
                hash_progress = progress("(1/2)  Calculating block hashes:")
77
                upload_progress = progress("(2/2)  Uploading missing blocks:")
78
                self.pithos_client.create_object(remote_path, f, size,
79
                                                hash_progress, upload_progress)
80
                self.uploaded_object = "pithos://%s/%s/%s" % \
81
                                (self.account, self.container, remote_path)
82
            except ClientError as e:
86
        try:
87
            self.pithos_client.create_container(self.container)
88
        except ClientError as e:
89
            if e.status != 202:  # Ignore container already exists errors
83 90
                raise FatalError("Pithos client: %d %s" % \
84
                                                        (e.status, e.message))
91
                                                    (e.status, e.message))
92
        try:
93
            hash_cb = progress(hp) if hp is not None else None
94
            upload_cb = progress(up) if up is not None else None
95
            self.pithos_client.create_object(remote_path, file_obj, size,
96
                                                            hash_cb, upload_cb)
97
            return "pithos://%s/%s/%s" % \
98
                            (self.account, self.container, remote_path)
99
        except ClientError as e:
100
            raise FatalError("Pithos client: %d %s" % (e.status, e.message))
85 101

  
86
    def register(self, metadata):
87
        pass
102
    def register(self, name, location, metadata):
103
        params = {'is_public':'true', 'disk_format':'diskdump'}
104
        try:
105
            self.image_client.register(name, location, params, metadata)
106
        except ClientError as e:
107
            raise FatalError("Image client: %d %s" % (e.status, e.message))
88 108

  
89 109
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
b/image_creator/main.py
43 43
import sys
44 44
import os
45 45
import optparse
46
import StringIO
46 47

  
47 48
dd = get_command('dd')
48 49

  
......
194 195
        size = options.shrink and dev.shrink() or dev.size()
195 196
        metadata['SIZE'] = str(size // 2 ** 20)
196 197

  
197
        #Calculating MD5sum
198
        output("Calculating md5sum...", False)
199 198
        checksum = md5(snapshot, size)
200
        success(checksum)
201
        output()
202 199

  
203
        if options.outfile is not None:
204
            f = open('%s.%s' % (options.outfile, 'meta'), 'w')
205
            try:
206
                for key in metadata.keys():
207
                    f.write("%s=%s\n" % (key, metadata[key]))
208
            finally:
209
                f.close()
200
        metastring = "\n".join(
201
            ['%s=%s' % (key, value) for (key, value) in metadata.items()])
210 202

  
203
        if options.outfile is not None:
211 204
            dev.dump(options.outfile)
212 205

  
206
            output('Dumping metadata file...', False)
207
            with open('%s.%s' % (options.outfile, 'meta'), 'w') as f:
208
                    f.write(metastring)
209
            success('done')
210

  
211
            output('Dumping md5sum file...', False)
212
            with open('%s.%s' % (options.outfile, 'md5sum'), 'w') as f:
213
                f.write('%s %s'% (checksum, os.path.basename(options.outfile)))
214
            success('done')
215

  
213 216
        # Destroy the device. We only need the snapshot from now on
214 217
        disk.destroy_device(dev)
215 218

  
219
        output()
220

  
221
        uploaded_obj = ""
216 222
        if options.upload:
217 223
            output("Uploading image to pithos:")
218 224
            kamaki = Kamaki(options.account, options.token)
219
            kamaki.upload(snapshot, size, options.upload)
225
            with open(snapshot) as f:
226
                uploaded_obj = kamaki.upload(f, size, options.upload,
227
                                "(1/4)  Calculating block hashes",
228
                                "(2/4)  Uploading missing blocks")
229

  
230
            output("(3/4)  Uploading metadata file...", False)
231
            kamaki.upload(StringIO.StringIO(metastring), size=len(metastring),
232
                                remote_path="%s.%s" % (options.upload, 'meta'))
233
            success('done')
234
            output("(4/4)  Uploading md5sum file...", False)
235
            md5sumstr = '%s %s' % (checksum, os.path.basename(options.upload))
236
            kamaki.upload(StringIO.StringIO(md5sumstr), size=len(md5sumstr),
237
                            remote_path="%s.%s" % (options.upload, 'md5sum'))
238
            success('done')
239
            output()
240

  
241
        if options.register:
242
            output('Registing image to ~okeanos...')
243
            kamaki.register(options.register, uploaded_obj, metadata)
244
            output('done')
245
            output()
220 246

  
221 247
    finally:
222 248
        output('cleaning up...')
223 249
        disk.cleanup()
224 250

  
251
    success("snf-image-creator exited without errors")
252

  
225 253
    return 0
226 254

  
227 255

  
b/image_creator/util.py
34 34
import sys
35 35
import pbs
36 36
import hashlib
37
from clint.textui import colored, progress as uiprogress
37
from clint.textui import colored
38
from progress.bar import Bar
38 39

  
39 40

  
40 41
class FatalError(Exception):
......
85 86
            sys.stdout.flush()
86 87

  
87 88

  
88
def progress(message=''):
89
def progress(message='', bar_type="default"):
89 90

  
90
    PROGRESS_LENGTH = 32
91
    MESSAGE_LENGTH = 32
91
    MESSAGE_LENGTH = 30
92
    
93
    suffix={'default':'%(index)d/%(max)d',
94
        'percent':'%(percent)d%%',
95
        'b':'%(index)d/%(max)d B',
96
        'kb':'%(index)d/%(max)d KB',
97
        'mb':'%(index)d/%(max)d MB'}
92 98

  
93
    def progress_generator(n=100):
94
        position = 0
95
        msg = message.ljust(MESSAGE_LENGTH)
96
        for i in uiprogress.bar(range(n), msg, PROGRESS_LENGTH, silent):
97
            if i < position:
98
                continue
99
            position = yield
100
        yield  # suppress the StopIteration exception
101
    return progress_generator
99
    return Bar(message=message.ljust(MESSAGE_LENGTH), fill='#', \
100
                                                    suffix=suffix[bar_type])
102 101

  
102
def md5(filename, size):
103 103

  
104
def md5(filename, size, progress=None):
105

  
106
    BLOCKSIZE = 2 ^ 22  # 4MB
104
    BLOCKSIZE = 2 ** 22  # 4MB
107 105

  
106
    progressbar = progress("Calculating md5sum:", 'mb')
107
    progressbar.max = (size // (2 ** 20))
108 108
    md5 = hashlib.md5()
109 109
    with open(filename, "r") as src:
110 110
        left = size
......
113 113
            data = src.read(length)
114 114
            md5.update(data)
115 115
            left -= length
116
            progressbar.goto((size - left) // (2 ** 20))
117
    
118
    checksum = md5.hexdigest()
119
    output("\rCalculating md5sum...\033[K", False)
120
    success(checksum)
116 121

  
117
    return md5.hexdigest()
122
    return checksum
118 123

  
119 124
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
b/setup.py
47 47
    license='BSD',
48 48
    packages=['image_creator', 'image_creator.os_type'],
49 49
    include_package_data=True,
50
    install_requires=['pbs', 'clint', 'pysendfile'],
50
    install_requires=['pbs', 'clint', 'progress','pysendfile'],
51 51
    entry_points={
52 52
        'console_scripts': ['snf-image-creator = image_creator.main:main']
53 53
    }

Also available in: Unified diff