Issue #134: use requests instead of urllib

This commit is contained in:
Michael Penkov
2017-10-28 16:11:58 +02:00
parent 66bda718c5
commit eabc79c8ec
2 changed files with 22 additions and 17 deletions

View File

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

View File

@ -4,7 +4,7 @@ import os
import runpy import runpy
import sys import sys
from contextlib import redirect_stdout, redirect_stderr from contextlib import redirect_stdout, redirect_stderr
from unittest.mock import patch from unittest.mock import patch, Mock
import pkg_resources import pkg_resources
import pytest import pytest
@ -155,13 +155,21 @@ def test_version(runner):
@all_runners @all_runners
def test_download_works_for_examples(runner): 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 = io.BytesIO()
fout.close = lambda: None fout.close = lambda: None
expected_bytes = b'hello world'
with patch('bonobo.commands.download._open_url') as mock_open_url, \ with patch('bonobo.commands.download._open_url') as mock_open_url, \
patch('bonobo.commands.download.open') as mock_open: 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 mock_open.return_value = fout
runner('download', 'examples/datasets/coffeeshops.txt') runner('download', 'examples/datasets/coffeeshops.txt')
expected_url = EXAMPLES_BASE_URL + 'datasets/coffeeshops.txt' expected_url = EXAMPLES_BASE_URL + 'datasets/coffeeshops.txt'