Merge pull request #44 from hbmartin/request-head

pytube/request.py

@@ -1,17 +1,22 @@
 # -*- coding: utf-8 -*-
 
 """Implements a simple wrapper around urlopen."""
-from typing import Any, Iterable, Dict
+from typing import Any, Iterable, Dict, Optional
 from urllib.request import Request
 from urllib.request import urlopen
 
 
-def _execute_request(url: str) -> Any:
-    if not url.lower().startswith("http"):
+def _execute_request(
+    url: str, method: Optional[str] = None, headers: Optional[Dict[str, str]] = None
+) -> Any:
+    base_headers = {"User-Agent": "Mozilla/5.0"}
+    if headers:
+        base_headers.update(headers)
+    if url.lower().startswith("http"):
+        request = Request(url, headers=base_headers, method=method)
+    else:
         raise ValueError
-    return urlopen(
-        Request(url, headers={"User-Agent": "Mozilla/5.0", "Range": "bytes=0-"})
-    )  # nosec
+    return urlopen(request)  # nosec
 
 
 def get(url) -> str:
@@ -32,7 +37,7 @@ def stream(url: str, chunk_size: int = 8192) -> Iterable[bytes]:
     :param int chunk_size: The size in bytes of each chunk. Defaults to 8*1024
     :rtype: Iterable[bytes]
     """
-    response = _execute_request(url)
+    response = _execute_request(url, headers={"Range": "bytes=0-"})
     while True:
         buf = response.read(chunk_size)
         if not buf:
@@ -40,7 +45,7 @@ def stream(url: str, chunk_size: int = 8192) -> Iterable[bytes]:
         yield buf
 
 
-def headers(url: str) -> Dict:
+def head(url: str) -> Dict:
     """Fetch headers returned http GET request.
 
     :param str url:
@@ -49,4 +54,5 @@ def headers(url: str) -> Dict:
     :returns:
         dictionary of lowercase headers
     """
-    return {k.lower(): v for k, v in _execute_request(url).info().items()}
+    response_headers = _execute_request(url, method="HEAD").info()
+    return {k.lower(): v for k, v in response_headers.items()}

pytube/streams.py

@@ -140,7 +140,7 @@ class Stream:
             Filesize (in bytes) of the stream.
         """
         if self._filesize is None:
-            headers = request.headers(self.url)
+            headers = request.head(self.url)
             self._filesize = int(headers["content-length"])
         return self._filesize
 

tests/test_request.py

@@ -30,7 +30,7 @@ def test_headers(mock_urlopen):
     response = mock.Mock()
     response.info.return_value = {"content-length": "16384"}
     mock_urlopen.return_value = response
-    response = request.headers("http://fakeassurl.gov")
+    response = request.head("http://fakeassurl.gov")
     assert response == {"content-length": "16384"}
 
 

tests/test_streams.py

@@ -10,8 +10,8 @@ from pytube import Stream, streams
 
 
 def test_filesize(cipher_signature, mocker):
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "6796391"}
+    mocker.patch.object(request, "head")
+    request.head.return_value = {"content-length": "6796391"}
     assert cipher_signature.streams.first().filesize == 6796391
 
 
@@ -87,8 +87,8 @@ def test_views(cipher_signature):
 
 
 def test_download(cipher_signature, mocker):
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "head")
+    request.head.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):
@@ -97,8 +97,8 @@ def test_download(cipher_signature, mocker):
 
 
 def test_download_with_prefix(cipher_signature, mocker):
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "head")
+    request.head.return_value = {"content-length": "16384"}
     mocker.patch.object(request, "stream")
     request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])
     streams.target_directory = MagicMock(return_value="/target")
@@ -109,8 +109,8 @@ def test_download_with_prefix(cipher_signature, mocker):
 
 
 def test_download_with_filename(cipher_signature, mocker):
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "head")
+    request.head.return_value = {"content-length": "16384"}
     mocker.patch.object(request, "stream")
     request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])
     streams.target_directory = MagicMock(return_value="/target")
@@ -121,8 +121,8 @@ def test_download_with_filename(cipher_signature, mocker):
 
 
 def test_download_with_existing(cipher_signature, mocker):
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "head")
+    request.head.return_value = {"content-length": "16384"}
     mocker.patch.object(request, "stream")
     streams.target_directory = MagicMock(return_value="/target")
     mocker.patch.object(os.path, "isfile")
@@ -137,8 +137,8 @@ def test_download_with_existing(cipher_signature, mocker):
 
 
 def test_download_with_existing_no_skip(cipher_signature, mocker):
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "head")
+    request.head.return_value = {"content-length": "16384"}
     mocker.patch.object(request, "stream")
     request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])
     streams.target_directory = MagicMock(return_value="/target")
@@ -167,8 +167,8 @@ def test_on_progress_hook(cipher_signature, mocker):
     callback_fn = mock.MagicMock()
     cipher_signature.register_on_progress_callback(callback_fn)
 
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "head")
+    request.head.return_value = {"content-length": "16384"}
     mocker.patch.object(request, "stream")
     request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])
 
@@ -186,8 +186,8 @@ def test_on_complete_hook(cipher_signature, mocker):
     callback_fn = mock.MagicMock()
     cipher_signature.register_on_complete_callback(callback_fn)
 
-    mocker.patch.object(request, "headers")
-    request.headers.return_value = {"content-length": "16384"}
+    mocker.patch.object(request, "head")
+    request.head.return_value = {"content-length": "16384"}
     mocker.patch.object(request, "stream")
     request.stream.return_value = iter([str(random.getrandbits(8 * 1024))])