Show a pretty progressbar when uploading and downloading an image.

Add a new module that contain generic wrapper for file and iterator, which
are used to wrap image to upload and the request body iterator in upload and
download cases repectively, to show and advance a pretty progress bar when this
laters are consumed, The progress bar is triggered by adding a --progress command
line argument to commands: image-create, image-download or image-update.

Change-Id: I2ba42fd0c58f4fa087adb568ec3f08246cae3759
bug fix: LP#1112309
blueprint: progressbar-when-uploading

glanceclient/common/http.py

@@ -14,6 +14,8 @@
 #    under the License.
 
 import copy
+import errno
+import hashlib
 import httplib
 import logging
 import posixpath
@@ -435,17 +437,55 @@ class VerifiedHTTPSConnection(HTTPSConnection):
 
 
 class ResponseBodyIterator(object):
-    """A class that acts as an iterator over an HTTP response."""
+    """
+    A class that acts as an iterator over an HTTP response.
+
+    This class will also check response body integrity when iterating over
+    the instance and if a checksum was supplied using `set_checksum` method,
+    else by default the class will not do any integrity check.
+    """
 
     def __init__(self, resp):
-        self.resp = resp
+        self._resp = resp
+        self._checksum = None
+        self._size = int(resp.getheader('content-length', 0))
+        self._end_reached = False
+
+    def set_checksum(self, checksum):
+        """
+        Set checksum to check against when iterating over this instance.
+
+        :raise: AttributeError if iterator is already consumed.
+        """
+        if self._end_reached:
+            raise AttributeError("Can't set checksum for an already consumed"
+                                 " iterator")
+        self._checksum = checksum
+
+    def __len__(self):
+        return int(self._size)
 
     def __iter__(self):
+        md5sum = hashlib.md5()
         while True:
-            yield self.next()
+            try:
+                chunk = self.next()
+            except StopIteration:
+                self._end_reached = True
+                # NOTE(mouad): Check image integrity when the end of response
+                # body is reached.
+                md5sum = md5sum.hexdigest()
+                if self._checksum is not None and md5sum != self._checksum:
+                    raise IOError(errno.EPIPE,
+                                  'Corrupted image. Checksum was %s '
+                                  'expected %s' % (md5sum, self._checksum))
+                raise
+            else:
+                yield chunk
+                md5sum.update(chunk)
 
     def next(self):
-        chunk = self.resp.read(CHUNKSIZE)
+        chunk = self._resp.read(CHUNKSIZE)
         if chunk:
             return chunk
         else:

glanceclient/common/progressbar.py

@@ -0,0 +1,87 @@
+# Copyright 2013 OpenStack LLC.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import sys
+
+
+class _ProgressBarBase(object):
+    """
+    Base abstract class used by specific class wrapper to show a progress bar
+    when the wrapped object are consumed.
+
+    :param wrapped: Object to wrap that hold data to be consumed.
+    :param totalsize: The total size of the data in the wrapped object.
+
+    :note: The progress will be displayed only if sys.stdout is a tty.
+    """
+
+    def __init__(self, wrapped, totalsize):
+        self._wrapped = wrapped
+        self._totalsize = float(totalsize)
+        self._show_progress = sys.stdout.isatty() and self._totalsize != 0
+        self._percent = 0
+
+    def _display_progress_bar(self, size_read):
+        if self._show_progress:
+            self._percent += size_read / self._totalsize
+            # Output something like this: [==========>             ] 49%
+            sys.stdout.write('\r[{0:<30}] {1:.0%}'.format(
+                '=' * int(round(self._percent * 29)) + '>', self._percent
+            ))
+            sys.stdout.flush()
+
+    def __getattr__(self, attr):
+        # Forward other attribute access to the wrapped object.
+        return getattr(self._wrapped, attr)
+
+
+class VerboseFileWrapper(_ProgressBarBase):
+    """
+    A file wrapper that show and advance a progress bar whenever file's read
+    method is called.
+    """
+
+    def read(self, *args, **kwargs):
+        data = self._wrapped.read(*args, **kwargs)
+        if data:
+            self._display_progress_bar(len(data))
+        else:
+            # Break to a new line from the progress bar for incoming output.
+            sys.stdout.write('\n')
+        return data
+
+
+class VerboseIteratorWrapper(_ProgressBarBase):
+    """
+    An iterator wrapper that show and advance a progress bar whenever
+    data is consumed from the iterator.
+
+    :note: Use only with iterator that yield strings.
+    """
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        try:
+            data = self._wrapped.next()
+            # NOTE(mouad): Assuming that data is a string b/c otherwise calling
+            # len function will not make any sense.
+            self._display_progress_bar(len(data))
+            return data
+        except StopIteration:
+            # Break to a new line from the progress bar for incoming output.
+            sys.stdout.write('\n')
+            raise

glanceclient/common/utils.py

@@ -14,7 +14,6 @@
 #    under the License.
 
 import errno
-import hashlib
 import os
 import sys
 import uuid
@@ -161,23 +160,6 @@ def save_image(data, path):
             image.close()
 
 
-def integrity_iter(iter, checksum):
-    """
-    Check image data integrity.
-
-    :raises: IOError
-    """
-    md5sum = hashlib.md5()
-    for chunk in iter:
-        yield chunk
-        md5sum.update(chunk)
-    md5sum = md5sum.hexdigest()
-    if md5sum != checksum:
-        raise IOError(errno.EPIPE,
-                      'Corrupt image download. Checksum was %s expected %s' %
-                      (md5sum, checksum))
-
-
 def make_size_human_readable(size):
     suffix = ['B', 'kB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB']
     base = 1024.0
@@ -214,3 +196,31 @@ def exception_to_str(exc):
             error = ("Caught '%(exception)s' exception." %
                      {"exception": exc.__class__.__name__})
     return strutils.safe_encode(error, errors='ignore')
+
+
+def get_file_size(file_obj):
+    """
+    Analyze file-like object and attempt to determine its size.
+
+    :param file_obj: file-like object.
+    :retval The file's size or None if it cannot be determined.
+    """
+    if hasattr(file_obj, 'seek') and hasattr(file_obj, 'tell'):
+        try:
+            curr = file_obj.tell()
+            file_obj.seek(0, os.SEEK_END)
+            size = file_obj.tell()
+            file_obj.seek(curr)
+            return size
+        except IOError, e:
+            if e.errno == errno.ESPIPE:
+                # Illegal seek. This means the file object
+                # is a pipe (e.g the user is trying
+                # to pipe image data to the client,
+                # echo testdata | bin/glance add blah...), or
+                # that file object is empty, or that a file-like
+                # object which doesn't support 'seek/tell' has
+                # been supplied.
+                return
+            else:
+                raise

glanceclient/v1/images.py

@@ -14,9 +14,7 @@
 #    under the License.
 
 import copy
-import errno
 import json
-import os
 import urllib
 
 from glanceclient.common import base
@@ -129,9 +127,8 @@ class ImageManager(base.Manager):
                                           % urllib.quote(image_id))
         checksum = resp.getheader('x-image-meta-checksum', None)
         if do_checksum and checksum is not None:
-            return utils.integrity_iter(body, checksum)
+            body.set_checksum(checksum)
-        else:
+        return body
-            return body
 
     def list(self, **kwargs):
         """Get a list of images.
@@ -228,35 +225,6 @@ class ImageManager(base.Manager):
         """Delete an image."""
         self._delete("/v1/images/%s" % base.getid(image))
 
-    def _get_file_size(self, obj):
-        """Analyze file-like object and attempt to determine its size.
-
-        :param obj: file-like object, typically redirected from stdin.
-        :retval The file's size or None if it cannot be determined.
-        """
-        # For large images, we need to supply the size of the
-        # image file. See LP Bugs #827660 and #845788.
-        if hasattr(obj, 'seek') and hasattr(obj, 'tell'):
-            try:
-                obj.seek(0, os.SEEK_END)
-                obj_size = obj.tell()
-                obj.seek(0)
-                return obj_size
-            except IOError as e:
-                if e.errno == errno.ESPIPE:
-                    # Illegal seek. This means the user is trying
-                    # to pipe image data to the client, e.g.
-                    # echo testdata | bin/glance add blah..., or
-                    # that stdin is empty, or that a file-like
-                    # object which doesn't support 'seek/tell' has
-                    # been supplied.
-                    return None
-                else:
-                    raise
-        else:
-            # Cannot determine size of input image
-            return None
-
     def create(self, **kwargs):
         """Create an image
 
@@ -264,7 +232,7 @@ class ImageManager(base.Manager):
         """
         image_data = kwargs.pop('data', None)
         if image_data is not None:
-            image_size = self._get_file_size(image_data)
+            image_size = utils.get_file_size(image_data)
             if image_size is not None:
                 kwargs.setdefault('size', image_size)
 
@@ -293,7 +261,7 @@ class ImageManager(base.Manager):
         """
         image_data = kwargs.pop('data', None)
         if image_data is not None:
-            image_size = self._get_file_size(image_data)
+            image_size = utils.get_file_size(image_data)
             if image_size is not None:
                 kwargs.setdefault('size', image_size)
 

glanceclient/v1/shell.py

@@ -25,6 +25,7 @@ else:
 
 from glanceclient import exc
 from glanceclient.common import utils
+from glanceclient.common import progressbar
 from glanceclient.openstack.common import strutils
 import glanceclient.v1.images
 
@@ -165,10 +166,14 @@ def do_image_show(gc, args):
                 'If this is not specified the image data will be '
                 'written to stdout.')
 @utils.arg('image', metavar='<IMAGE>', help='Name or ID of image to download.')
+@utils.arg('--progress', action='store_true', default=False,
+           help='Show download progress bar.')
 def do_image_download(gc, args):
     """Download a specific image."""
     image = utils.find_resource(gc.images, args.image)
     body = image.data()
+    if args.progress:
+        body = progressbar.VerboseIteratorWrapper(body, len(body))
     utils.save_image(body, args.file)
 
 
@@ -219,6 +224,8 @@ def do_image_download(gc, args):
                  "May be used multiple times."))
 @utils.arg('--human-readable', action='store_true', default=False,
            help='Print image size in a human-friendly format.')
+@utils.arg('--progress', action='store_true', default=False,
+           help='Show upload progress bar.')
 def do_image_create(gc, args):
     """Create a new image."""
     # Filter out None values
@@ -241,6 +248,12 @@ def do_image_create(gc, args):
 
     _set_data_field(fields, args)
 
+    if args.progress:
+        filesize = utils.get_file_size(fields['data'])
+        fields['data'] = progressbar.VerboseFileWrapper(
+            fields['data'], filesize
+        )
+
     image = gc.images.create(**fields)
     _image_show(image, args.human_readable)
 
@@ -287,6 +300,8 @@ def do_image_create(gc, args):
                  "those properties not referenced are preserved."))
 @utils.arg('--human-readable', action='store_true', default=False,
            help='Print image size in a human-friendly format.')
+@utils.arg('--progress', action='store_true', default=False,
+           help='Show upload progress bar.')
 def do_image_update(gc, args):
     """Update a specific image."""
     # Filter out None values
@@ -311,6 +326,12 @@ def do_image_update(gc, args):
     if image.status == 'queued':
         _set_data_field(fields, args)
 
+        if args.progress:
+            filesize = utils.get_file_size(fields['data'])
+            fields['data'] = progressbar.VerboseFileWrapper(
+                fields['data'], filesize
+            )
+
     image = gc.images.update(image, purge_props=args.purge_props, **fields)
     _image_show(image, args.human_readable)
 

glanceclient/v2/images.py

@@ -84,9 +84,8 @@ class Controller(object):
         resp, body = self.http_client.raw_request('GET', url)
         checksum = resp.getheader('content-md5', None)
         if do_checksum and checksum is not None:
-            return utils.integrity_iter(body, checksum)
+            body.set_checksum(checksum)
-        else:
+        return body
-            return body
 
     def upload(self, image_id, image_data):
         """
@@ -114,7 +113,7 @@ class Controller(object):
             try:
                 setattr(image, key, value)
             except warlock.InvalidOperation, e:
-                raise TypeError(unicode(message))
+                raise TypeError(utils.exception_to_str(e))
 
         resp, body = self.http_client.json_request('POST', url, body=image)
         #NOTE(esheffield): remove 'self' for now until we have an elegant

glanceclient/v2/shell.py

@@ -13,6 +13,7 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+from glanceclient.common import progressbar
 from glanceclient.common import utils
 from glanceclient import exc
 
@@ -125,9 +126,13 @@ def do_explain(gc, args):
                 'If this is not specified the image data will be '
                 'written to stdout.')
 @utils.arg('id', metavar='<IMAGE_ID>', help='ID of image to download.')
+@utils.arg('--progress', action='store_true', default=False,
+           help='Show download progress bar.')
 def do_image_download(gc, args):
     """Download a specific image."""
     body = gc.images.data(args.id)
+    if args.progress:
+        body = progressbar.VerboseIteratorWrapper(body, len(body))
     utils.save_image(body, args.file)
 
 

tests/test_http.py

@@ -13,6 +13,7 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+import errno
 import httplib
 import socket
 import StringIO
@@ -241,8 +242,42 @@ class TestHostResolutionError(testtools.TestCase):
 
 
 class TestResponseBodyIterator(testtools.TestCase):
+
     def test_iter_default_chunk_size_64k(self):
         resp = utils.FakeResponse({}, StringIO.StringIO('X' * 98304))
         iterator = http.ResponseBodyIterator(resp)
         chunks = list(iterator)
         self.assertEqual(chunks, ['X' * 65536, 'X' * 32768])
+
+    def test_integrity_check_with_correct_checksum(self):
+        resp = utils.FakeResponse({}, StringIO.StringIO('CCC'))
+        body = http.ResponseBodyIterator(resp)
+        body.set_checksum('defb99e69a9f1f6e06f15006b1f166ae')
+        list(body)
+
+    def test_integrity_check_with_wrong_checksum(self):
+        resp = utils.FakeResponse({}, StringIO.StringIO('BB'))
+        body = http.ResponseBodyIterator(resp)
+        body.set_checksum('wrong')
+        try:
+            list(body)
+            self.fail('integrity checked passed with wrong checksum')
+        except IOError as e:
+            self.assertEqual(errno.EPIPE, e.errno)
+
+    def test_set_checksum_in_consumed_iterator(self):
+        resp = utils.FakeResponse({}, StringIO.StringIO('CCC'))
+        body = http.ResponseBodyIterator(resp)
+        list(body)
+        # Setting checksum for an already consumed iterator should raise an
+        # AttributeError.
+        self.assertRaises(
+            AttributeError, body.set_checksum,
+            'defb99e69a9f1f6e06f15006b1f166ae')
+
+    def test_body_size(self):
+        size = 1000000007
+        resp = utils.FakeResponse(
+            {'content-length': str(size)}, StringIO.StringIO('BB'))
+        body = http.ResponseBodyIterator(resp)
+        self.assertEqual(len(body), size)

tests/test_progressbar.py

@@ -0,0 +1,59 @@
+# Copyright 2013 OpenStack LLC.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import sys
+import StringIO
+
+import testtools
+
+from glanceclient.common import progressbar
+from tests import utils as test_utils
+
+
+class TestProgressBarWrapper(testtools.TestCase):
+
+    def test_iter_iterator_display_progress_bar(self):
+        size = 100
+        iterator = iter('X' * 100)
+        saved_stdout = sys.stdout
+        try:
+            sys.stdout = output = test_utils.FakeTTYStdout()
+            # Consume iterator.
+            data = list(progressbar.VerboseIteratorWrapper(iterator, size))
+            self.assertEqual(data, ['X'] * 100)
+            self.assertEqual(
+                output.getvalue().strip(),
+                '[%s>] 100%%' % ('=' * 29)
+            )
+        finally:
+            sys.stdout = saved_stdout
+
+    def test_iter_file_display_progress_bar(self):
+        size = 98304
+        file_obj = StringIO.StringIO('X' * size)
+        saved_stdout = sys.stdout
+        try:
+            sys.stdout = output = test_utils.FakeTTYStdout()
+            file_obj = progressbar.VerboseFileWrapper(file_obj, size)
+            chunksize = 1024
+            chunk = file_obj.read(chunksize)
+            while chunk:
+                chunk = file_obj.read(chunksize)
+            self.assertEqual(
+                output.getvalue().strip(),
+                '[%s>] 100%%' % ('=' * 29)
+            )
+        finally:
+            sys.stdout = saved_stdout

tests/test_utils.py

@@ -13,39 +13,16 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
-import errno
-import testtools
 import sys
 import StringIO
 
+import testtools
+
 from glanceclient.common import utils
 
 
 class TestUtils(testtools.TestCase):
 
-    def test_integrity_iter_without_checksum(self):
-        try:
-            data = ''.join([f for f in utils.integrity_iter('A', None)])
-            self.fail('integrity checked passed without checksum.')
-        except IOError as e:
-            self.assertEqual(errno.EPIPE, e.errno)
-            msg = 'was 7fc56270e7a70fa81a5935b72eacbe29 expected None'
-            self.assertTrue(msg in str(e))
-
-    def test_integrity_iter_with_wrong_checksum(self):
-        try:
-            data = ''.join([f for f in utils.integrity_iter('BB', 'wrong')])
-            self.fail('integrity checked passed with wrong checksum')
-        except IOError as e:
-            self.assertEqual(errno.EPIPE, e.errno)
-            msg = 'was 9d3d9048db16a7eee539e93e3618cbe7 expected wrong'
-            self.assertTrue('expected wrong' in str(e))
-
-    def test_integrity_iter_with_checksum(self):
-        fixture = 'CCC'
-        checksum = 'defb99e69a9f1f6e06f15006b1f166ae'
-        data = ''.join([f for f in utils.integrity_iter(fixture, checksum)])
-
     def test_make_size_human_readable(self):
         self.assertEqual("106B", utils.make_size_human_readable(106))
         self.assertEqual("1000kB", utils.make_size_human_readable(1024000))
@@ -53,6 +30,27 @@ class TestUtils(testtools.TestCase):
         self.assertEqual("1.4GB", utils.make_size_human_readable(1476395008))
         self.assertEqual("9.3MB", utils.make_size_human_readable(9761280))
 
+    def test_get_new_file_size(self):
+        size = 98304
+        file_obj = StringIO.StringIO('X' * size)
+        try:
+            self.assertEqual(utils.get_file_size(file_obj), size)
+            # Check that get_file_size didn't change original file position.
+            self.assertEqual(file_obj.tell(), 0)
+        finally:
+            file_obj.close()
+
+    def test_get_consumed_file_size(self):
+        size, consumed = 98304, 304
+        file_obj = StringIO.StringIO('X' * size)
+        file_obj.seek(consumed)
+        try:
+            self.assertEqual(utils.get_file_size(file_obj), size)
+            # Check that get_file_size didn't change original file position.
+            self.assertEqual(file_obj.tell(), consumed)
+        finally:
+            file_obj.close()
+
     def test_prettytable(self):
         class Struct:
             def __init__(self, **entries):

tests/utils.py

@@ -33,8 +33,9 @@ class FakeAPI(object):
 
     def raw_request(self, *args, **kwargs):
         fixture = self._request(*args, **kwargs)
-        body_iter = http.ResponseBodyIterator(StringIO.StringIO(fixture[1]))
+        resp = FakeResponse(fixture[0], StringIO.StringIO(fixture[1]))
-        return FakeResponse(fixture[0]), body_iter
+        body_iter = http.ResponseBodyIterator(resp)
+        return resp, body_iter
 
     def json_request(self, *args, **kwargs):
         fixture = self._request(*args, **kwargs)
@@ -95,3 +96,17 @@ class TestResponse(requests.Response):
     @property
     def text(self):
         return self._text
+
+
+class FakeTTYStdout(StringIO.StringIO):
+    """A Fake stdout that try to emulate a TTY device as much as possible."""
+
+    def isatty(self):
+        return True
+
+    def write(self, data):
+        # When a CR (carriage return) is found reset file.
+        if data.startswith('\r'):
+            self.seek(0)
+            data = data[1:]
+        return StringIO.StringIO.write(self, data)

tests/v1/test_shell.py

@@ -365,7 +365,8 @@ class ShellStdinHandlingTests(testtools.TestCase):
                 property={},
                 purge_props=False,
                 human_readable=False,
-                file=None
+                file=None,
+                progress=False
             )
         )
 

tests/v2/test_shell_v2.py

@@ -15,11 +15,16 @@
 #    under the License.
 # vim: tabstop=4 shiftwidth=4 softtabstop=4
 
+import StringIO
+
 import mock
 import testtools
 
+from glanceclient.common import http
+from glanceclient.common import progressbar
 from glanceclient.common import utils
 from glanceclient.v2 import shell as test_shell
+from tests import utils as test_utils
 
 
 class ShellV2Test(testtools.TestCase):
@@ -109,15 +114,32 @@ class ShellV2Test(testtools.TestCase):
             self.gc.schemas.get.assert_called_once_with('test')
 
     def test_image_download(self):
-        args = self._make_args({'id': 'pass', 'file': 'test'})
+        args = self._make_args(
+            {'id': 'pass', 'file': 'test', 'progress': False})
 
         with mock.patch.object(self.gc.images, 'data') as mocked_data:
-            mocked_data.return_value = 'test_passed'
+            resp = test_utils.FakeResponse({}, StringIO.StringIO('CCC'))
-
+            ret = mocked_data.return_value = http.ResponseBodyIterator(resp)
             test_shell.do_image_download(self.gc, args)
 
             mocked_data.assert_called_once_with('pass')
-            utils.save_image.assert_called_once_with('test_passed', 'test')
+            utils.save_image.assert_called_once_with(ret, 'test')
+
+    def test_image_download_with_progressbar(self):
+        args = self._make_args(
+            {'id': 'pass', 'file': 'test', 'progress': True})
+
+        with mock.patch.object(self.gc.images, 'data') as mocked_data:
+            resp = test_utils.FakeResponse({}, StringIO.StringIO('CCC'))
+            mocked_data.return_value = http.ResponseBodyIterator(resp)
+            test_shell.do_image_download(self.gc, args)
+
+            mocked_data.assert_called_once_with('pass')
+            utils.save_image.assert_called_once_with(mock.ANY, 'test')
+            self.assertIsInstance(
+                utils.save_image.call_args[0][0],
+                progressbar.VerboseIteratorWrapper
+            )
 
     def test_do_image_delete(self):
         args = self._make_args({'id': 'pass', 'file': 'test'})