Issue #134: use requests instead of urllib

bonobo/commands/download.py

@@ -1,6 +1,7 @@
 import io
 import re
-import urllib.request
+
+import requests
 
 import bonobo
 
@@ -8,21 +9,17 @@ EXAMPLES_BASE_URL = 'https://raw.githubusercontent.com/python-bonobo/bonobo/mast
 """The URL to our git repository, in raw mode."""
 
 
-def _save_stream(fin, fout):
-    """Read the input stream and write it to the output stream block-by-block."""
-    while True:
-        data = fin.read(io.DEFAULT_BUFFER_SIZE)
-        if data:
-            fout.write(data)
-        else:
-            break
+def _write_response(response, fout):
+    """Read the response and write it to the output stream in chunks."""
+    for chunk in response.iter_content(io.DEFAULT_BUFFER_SIZE):
+        fout.write(chunk)
 
 
 def _open_url(url):
     """Open a HTTP connection to the URL and return a file-like object."""
-    response = urllib.request.urlopen(url)
-    if response.getcode() != 200:
-        raise IOError('unable to download {}, HTTP {}'.format(url, response.getcode()))
+    response = requests.get(url, stream=True)
+    if response.status_code != 200:
+        raise IOError('unable to download {}, HTTP {}'.format(url, response.status_code))
     return response
 
 
@@ -32,9 +29,9 @@ def execute(path, *args, **kwargs):
         raise ValueError('download command currently supports examples only')
     examples_path = re.sub('^examples/', '', path)
     output_path = bonobo.get_examples_path(examples_path)
-    fin = _open_url(EXAMPLES_BASE_URL + examples_path)
+    response = _open_url(EXAMPLES_BASE_URL + examples_path)
     with open(output_path, 'wb') as fout:
-        _save_stream(fin, fout)
+        _write_response(response, fout)
     print('saved to {}'.format(output_path))
 
 

tests/test_commands.py

@@ -4,7 +4,7 @@ import os
 import runpy
 import sys
 from contextlib import redirect_stdout, redirect_stderr
-from unittest.mock import patch
+from unittest.mock import patch, Mock
 
 import pkg_resources
 import pytest
@@ -155,13 +155,21 @@ def test_version(runner):
 
 @all_runners
 def test_download_works_for_examples(runner):
+    expected_bytes = b'hello world'
+
+    class MockResponse(object):
+        def __init__(self):
+            self.status_code = 200
+
+        def iter_content(self, *args, **kwargs):
+            return [expected_bytes]
+
     fout = io.BytesIO()
     fout.close = lambda: None
 
-    expected_bytes = b'hello world'
     with patch('bonobo.commands.download._open_url') as mock_open_url, \
             patch('bonobo.commands.download.open') as mock_open:
-        mock_open_url.return_value = io.BytesIO(expected_bytes)
+        mock_open_url.return_value = MockResponse()
         mock_open.return_value = fout
         runner('download', 'examples/datasets/coffeeshops.txt')
     expected_url = EXAMPLES_BASE_URL + 'datasets/coffeeshops.txt'