clean up request

pytube/request.py

@@ -1,39 +1,49 @@
 # -*- coding: utf-8 -*-
 """Implements a simple wrapper around urlopen."""
+from typing import Any, Iterable, Dict
 from urllib.request import Request
 from urllib.request import urlopen
 
 
-def get(url, headers=False, streaming=False, chunk_size=8192):
+def _execute_request(url:str) -> Any:
+    return urlopen(Request(url, headers={"User-Agent": "Mozilla/5.0"}))
+
+
+def get(url) -> str:
     """Send an http GET request.
 
     :param str url:
         The URL to perform the GET request for.
-    :param bool headers:
-        Only return the http headers.
-    :param bool streaming:
-        Returns the response body in chunks via a generator.
-    :param int chunk_size:
-        The size in bytes of each chunk. Defaults to 8*1024
+    :rtype: str
+    :returns:
+        UTF-8 encoded string of response
     """
+    return _execute_request(url).read().decode("utf-8")
 
-    req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
-    response = urlopen(req)
-
-    if streaming:
-        return stream_response(response, chunk_size)
 
-    if headers:
-        # https://github.com/nficano/pytube/issues/160
-        return {k.lower(): v for k, v in response.info().items()}
-
-    return response.read().decode("utf-8")
-
-
-def stream_response(response, chunk_size=8 * 1024):
-    """Read the response in chunks."""
+def stream(url: str, chunk_size:int = 8192) -> Iterable[bytes]:
+    """Read the response in chunks.
+    :param str url:
+        The URL to perform the GET request for.
+    :param int chunk_size:
+        The size in bytes of each chunk. Defaults to 8*1024
+    :rtype: Iterable[bytes]
+    """
+    response = _execute_request(url)
     while True:
         buf = response.read(chunk_size)
         if not buf:
             break
         yield buf
+
+
+def headers(url: str) -> Dict:
+    """Fetch headers returned http GET request.
+
+    :param str url:
+        The URL to perform the GET request for.
+    :rtype: dict
+    :returns:
+        dictionary of lowercase headers
+    """
+    return {k.lower(): v for k, v in _execute_request(url).info().items()}

pytube/streams.py

@@ -160,7 +160,7 @@ class Stream:
             Filesize (in bytes) of the stream.
         """
         if self._filesize is None:
-            headers = request.get(self.url, headers=True)
+            headers = request.headers(self.url)
             self._filesize = int(headers["content-length"])
         return self._filesize
 
@@ -239,7 +239,7 @@ class Stream:
         )
 
         with open(file_path, "wb") as fh:
-            for chunk in request.get(self.url, streaming=True):
+            for chunk in request.stream(self.url):
                 # reduce the (bytes) remainder by the length of the chunk.
                 bytes_remaining -= len(chunk)
                 # send to the on_progress callback.
@@ -258,7 +258,7 @@ class Stream:
             "downloading (%s total bytes) file to BytesIO buffer", self.filesize,
         )
 
-        for chunk in request.get(self.url, streaming=True):
+        for chunk in request.stream(self.url):
             # reduce the (bytes) remainder by the length of the chunk.
             bytes_remaining -= len(chunk)
             # send to the on_progress callback.

tests/test_request.py

@@ -7,7 +7,7 @@ from pytube import request
 
 
 @mock.patch("pytube.request.urlopen")
-def test_get_streaming(mock_urlopen):
+def test_streaming(mock_urlopen):
     fake_stream_binary = [
         iter(os.urandom(8 * 1024)),
         iter(os.urandom(8 * 1024)),
@@ -17,18 +17,18 @@ def test_get_streaming(mock_urlopen):
     response = mock.Mock()
     response.read.side_effect = fake_stream_binary
     mock_urlopen.return_value = response
-    response = request.get("http://fakeassurl.gov", streaming=True)
+    response = request.stream("http://fakeassurl.gov")
     call_count = len(list(response))
 
     assert call_count == 3
 
 
 @mock.patch("pytube.request.urlopen")
-def test_get_headers(mock_urlopen):
+def test_headers(mock_urlopen):
     response = mock.Mock()
     response.info.return_value = {"content-length": "16384"}
     mock_urlopen.return_value = response
-    response = request.get("http://fakeassurl.gov", headers=True)
+    response = request.headers("http://fakeassurl.gov")
     assert response == {"content-length": "16384"}
 
 

tests/test_streams.py

@@ -8,8 +8,8 @@ from pytube import Stream
 
 
 def test_filesize(cipher_signature, mocker):
-    mocker.patch.object(request, "get")
-    request.get.return_value = {"content-length": "6796391"}
+    mocker.patch.object(request, "headers")
+    request.headers.return_value = {"content-length": "6796391"}
     assert cipher_signature.streams.first().filesize == 6796391
 
 
@@ -36,12 +36,10 @@ def test_title(cipher_signature):
 
 
 def test_download(cipher_signature, mocker):
-    mocker.patch.object(request, "get")
-    request.get.side_effect = [
-        {"content-length": "16384"},
-        {"content-length": "16384"},
-        iter([str(random.getrandbits(8 * 1024))]),
-    ]
+    mocker.patch.object(request, "headers")
+    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "stream")
+    request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])
     with mock.patch("pytube.streams.open", mock.mock_open(), create=True):
         stream = cipher_signature.streams.first()
         stream.download()
@@ -61,12 +59,11 @@ def test_on_progress_hook(cipher_signature, mocker):
     callback_fn = mock.MagicMock()
     cipher_signature.register_on_progress_callback(callback_fn)
 
-    mocker.patch.object(request, "get")
-    request.get.side_effect = [
-        {"content-length": "16384"},
-        {"content-length": "16384"},
-        iter([str(random.getrandbits(8 * 1024))]),
-    ]
+    mocker.patch.object(request, "headers")
+    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "stream")
+    request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])
+
     with mock.patch("pytube.streams.open", mock.mock_open(), create=True):
         stream = cipher_signature.streams.first()
         stream.download()
@@ -81,12 +78,11 @@ def test_on_complete_hook(cipher_signature, mocker):
     callback_fn = mock.MagicMock()
     cipher_signature.register_on_complete_callback(callback_fn)
 
-    mocker.patch.object(request, "get")
-    request.get.side_effect = [
-        {"content-length": "16384"},
-        {"content-length": "16384"},
-        iter([str(random.getrandbits(8 * 1024))]),
-    ]
+    mocker.patch.object(request, "headers")
+    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "stream")
+    request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])
+
     with mock.patch("pytube.streams.open", mock.mock_open(), create=True):
         stream = cipher_signature.streams.first()
         stream.download()