Add progressbar support to most download actions (#615)

LP: #1597387

Signed-off-by: Sergio Schvezov <sergio.schvezov@ubuntu.com>

snapcraft/internal/indicators.py

@@ -0,0 +1,52 @@
+# -*- Mode:Python; indent-tabs-mode:nil; tab-width:4 -*-
+#
+# Copyright (C) 2016 Canonical Ltd
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 3 as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import os
+
+from progressbar import (
+    AnimatedMarker,
+    Bar,
+    Percentage,
+    ProgressBar,
+    UnknownLength,
+)
+
+
+def download_requests_stream(request_stream, destination, message=None):
+    """This is a facility to download a request with nice progress bars."""
+    if not message:
+        message = 'Downloading {!r}'.format(os.path.basename(destination))
+
+    total_length = int(request_stream.headers.get('Content-Length', '0'))
+    if total_length:
+        progress_bar = ProgressBar(
+            widgets=[message,
+                     Bar(marker='=', left='[', right=']'),
+                     ' ', Percentage()],
+            maxval=total_length)
+    else:
+        progress_bar = ProgressBar(
+            widgets=[message, AnimatedMarker()],
+            maxval=UnknownLength)
+
+    total_read = 0
+    progress_bar.start()
+    with open(destination, 'wb') as destination_file:
+        for buf in request_stream.iter_content(1024):
+            destination_file.write(buf)
+            total_read += len(buf)
+            progress_bar.update(total_read)
+    progress_bar.finish()

snapcraft/internal/parts.py

@@ -21,15 +21,9 @@
 
 import requests
 import yaml
-from progressbar import (
-    AnimatedMarker,
-    Bar,
-    Percentage,
-    ProgressBar,
-    UnknownLength,
-)
 from xdg import BaseDirectory
 
+from snapcraft.internal.indicators import download_requests_stream
 from snapcraft.internal.common import get_terminal_width
 
 
@@ -67,31 +61,10 @@ def execute(self):
             return
         self._request.raise_for_status()
 
-        self._download()
+        download_requests_stream(self._request, self.parts_yaml,
+                                 'Downloading parts list')
         self._save_headers()
 
-    def _download(self):
-        total_length = int(self._request.headers.get('Content-Length', '0'))
-        if total_length:
-            progress_bar = ProgressBar(
-                widgets=['Updating parts list',
-                         Bar(marker='=', left='[', right=']'),
-                         ' ', Percentage()],
-                maxval=total_length)
-        else:
-            progress_bar = ProgressBar(
-                widgets=['Updating parts list... ', AnimatedMarker()],
-                maxval=UnknownLength)
-
-        total_read = 0
-        progress_bar.start()
-        with open(self.parts_yaml, 'wb') as parts_file:
-            for buf in self._request.iter_content(1024):
-                parts_file.write(buf)
-                total_read += len(buf)
-                progress_bar.update(total_read)
-        progress_bar.finish()
-
     def _load_headers(self):
         if not os.path.exists(self._headers_yaml):
             return None

snapcraft/internal/sources.py

@@ -71,6 +71,7 @@
 import zipfile
 
 from snapcraft.internal import common
+from snapcraft.internal.indicators import download_requests_stream
 
 
 logging.getLogger('urllib3').setLevel(logging.CRITICAL)
@@ -103,15 +104,12 @@ def pull(self):
         self.provision(self.source_dir)
 
     def download(self):
-        req = requests.get(self.source, stream=True, allow_redirects=True)
-        if req.status_code is not 200:
-            raise EnvironmentError('unexpected http status code when '
-                                   'downloading {}'.format(req.status_code))
-
-        file = os.path.join(self.source_dir, os.path.basename(self.source))
-        with open(file, 'wb') as f:
-            for chunk in req.iter_content(1024):
-                f.write(chunk)
+        request = requests.get(self.source, stream=True, allow_redirects=True)
+        request.raise_for_status()
+
+        file_path = os.path.join(
+            self.source_dir, os.path.basename(self.source))
+        download_requests_stream(request, file_path)
 
 
 class Bazaar(Base):

snapcraft/storeapi/__init__.py

@@ -28,6 +28,7 @@
 
 import snapcraft
 from snapcraft import config
+from snapcraft.internal.indicators import download_requests_stream
 from snapcraft.storeapi import (
     _upload,
     constants,
@@ -191,12 +192,10 @@ def _download_snap(self, name, channel, arch, download_path,
                 name, download_path))
             return
         logger.info('Downloading {}'.format(name, download_path))
-        # FIXME: Check the status code ! -- vila 2016-05-04
-        download = self.cpi.get(download_url)
-        with open(download_path, 'wb') as f:
-            # FIXME: Cough, we may want to buffer here (and a progress bar
-            # would be nice) -- vila 2016-04-26
-            f.write(download.content)
+        request = self.cpi.get(download_url, stream=True)
+        request.raise_for_status()
+        download_requests_stream(request, download_path)
+
         if self._is_downloaded(download_path, expected_sha512):
             logger.info('Successfully downloaded {} at {}'.format(
                 name, download_path))
@@ -275,11 +274,12 @@ def search_package(self, snap_name, channel, arch):
         else:
             return embedded['clickindex:package'][0]
 
-    def get(self, url, headers=None, params=None):
+    def get(self, url, headers=None, params=None, stream=False):
         if headers is None:
             headers = {}
         headers.update({'Authorization': _macaroon_auth(self.conf)})
-        response = self.request('GET', url, headers=headers, params=params)
+        response = self.request('GET', url, stream=stream,
+                                headers=headers, params=params)
         return response
 
 

snapcraft/tests/__init__.py

@@ -62,6 +62,12 @@ def setUp(self):
         self.cpu_count.return_value = 2
         self.addCleanup(patcher.stop)
 
+        patcher = mock.patch(
+            'snapcraft.internal.indicators.ProgressBar',
+            new=SilentProgressBar)
+        patcher.start()
+        self.addCleanup(patcher.stop)
+
         # These are what we expect by default
         self.snap_dir = os.path.join(os.getcwd(), 'prime')
         self.stage_dir = os.path.join(os.getcwd(), 'stage')

snapcraft/tests/test_commands_define.py

@@ -27,9 +27,7 @@ class DefineCommandTestCase(tests.TestCase):
     def setUp(self):
         super().setUp()
         self.useFixture(fixture_setup.FakeParts())
-        with mock.patch('snapcraft.internal.parts.ProgressBar',
-                        new=tests.SilentProgressBar):
-            parts.update()
+        parts.update()
 
     @mock.patch('sys.stdout', new_callable=io.StringIO)
     def test_defining_a_part_that_exists(self, mock_stdout):

snapcraft/tests/test_commands_search.py

@@ -15,7 +15,6 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import logging
-from unittest import mock
 
 import fixtures
 
@@ -29,9 +28,7 @@ class SearchCommandTestCase(tests.TestCase):
     def setUp(self):
         super().setUp()
         self.useFixture(fixture_setup.FakeParts())
-        with mock.patch('snapcraft.internal.parts.ProgressBar',
-                        new=tests.SilentProgressBar):
-            parts.update()
+        parts.update()
 
     def test_searching_for_a_part_that_exists(self):
         fake_terminal = fixture_setup.FakeTerminal()

snapcraft/tests/test_commands_update.py

@@ -17,7 +17,6 @@
 
 import logging
 import os
-from unittest import mock
 
 import fixtures
 import yaml
@@ -32,11 +31,6 @@ class UpdateCommandTestCase(tests.TestCase):
     def setUp(self):
         super().setUp()
         self.useFixture(fixture_setup.FakeParts())
-        patcher = mock.patch(
-            'snapcraft.internal.parts.ProgressBar',
-            new=tests.SilentProgressBar)
-        patcher.start()
-        self.addCleanup(patcher.stop)
 
         self.parts_dir = os.path.join(BaseDirectory.xdg_data_home, 'snapcraft')
         self.parts_yaml = os.path.join(self.parts_dir, 'parts.yaml')

snapcraft/tests/test_yaml.py

@@ -93,11 +93,6 @@ def test_config_loads_with_different_encodings(
     @unittest.mock.patch('snapcraft.internal.yaml.Config.load_plugin')
     def test_config_composes_with_remote_parts(self, mock_loadPlugin):
         self.useFixture(fixture_setup.FakeParts())
-        patcher = unittest.mock.patch(
-            'snapcraft.internal.parts.ProgressBar',
-            new=tests.SilentProgressBar)
-        patcher.start()
-        self.addCleanup(patcher.stop)
 
         self.make_snapcraft_yaml("""name: test
 version: "1"
@@ -119,11 +114,6 @@ def test_config_composes_with_remote_parts(self, mock_loadPlugin):
 
     def test_config_composes_with_a_non_existent_remote_part(self):
         self.useFixture(fixture_setup.FakeParts())
-        patcher = unittest.mock.patch(
-            'snapcraft.internal.parts.ProgressBar',
-            new=tests.SilentProgressBar)
-        patcher.start()
-        self.addCleanup(patcher.stop)
 
         self.make_snapcraft_yaml("""name: test
 version: "1"
@@ -148,11 +138,6 @@ def test_config_composes_with_a_non_existent_remote_part(self):
 
     def test_config_after_is_an_undefined_part(self):
         self.useFixture(fixture_setup.FakeParts())
-        patcher = unittest.mock.patch(
-            'snapcraft.internal.parts.ProgressBar',
-            new=tests.SilentProgressBar)
-        patcher.start()
-        self.addCleanup(patcher.stop)
 
         self.make_snapcraft_yaml("""name: test
 version: "1"
@@ -179,11 +164,6 @@ def test_config_after_is_an_undefined_part(self):
     @unittest.mock.patch('snapcraft.internal.pluginhandler.load_plugin')
     def test_config_uses_remote_part_from_after(self, mock_load):
         self.useFixture(fixture_setup.FakeParts())
-        patcher = unittest.mock.patch(
-            'snapcraft.internal.parts.ProgressBar',
-            new=tests.SilentProgressBar)
-        patcher.start()
-        self.addCleanup(patcher.stop)
 
         self.make_snapcraft_yaml("""name: test
 version: "1"