Merge pull request #64 from hbmartin/range-tuning

better range stream defaults, avoid unnecessary head call

pytube/request.py

@@ -1,22 +1,26 @@
 # -*- coding: utf-8 -*-
 
 """Implements a simple wrapper around urlopen."""
+import logging
 from functools import lru_cache
-from typing import Any, Iterable, Dict, Optional
+from http.client import HTTPResponse
+from typing import Iterable, Dict, Optional
 from urllib.request import Request
 from urllib.request import urlopen
 
+logger = logging.getLogger(__name__)
+
 
 def _execute_request(
     url: str, method: Optional[str] = None, headers: Optional[Dict[str, str]] = None
-) -> Any:
+) -> HTTPResponse:
     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
+        raise ValueError("Invalid URL")
     return urlopen(request)  # nosec
 
 
@@ -33,20 +37,26 @@ def get(url) -> str:
 
 
 def stream(
-    url: str, chunk_size: int = 1024, range_size: int = 10485760
+    url: str, chunk_size: int = 4096, range_size: int = 9437184
 ) -> 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 1KB
-    :param int range_size: The size in bytes of each range request. Defaults to 10MB
+    :param int chunk_size: The size in bytes of each chunk. Defaults to 4KB
+    :param int range_size: The size in bytes of each range request. Defaults to 9MB
     :rtype: Iterable[bytes]
     """
-    file_size: int = filesize(url)
+    file_size: int = range_size  # fake filesize to start
     downloaded = 0
     while downloaded < file_size:
         stop_pos = min(downloaded + range_size, file_size) - 1
         range_header = f"bytes={downloaded}-{stop_pos}"
         response = _execute_request(url, method="GET", headers={"Range": range_header})
+        if file_size == range_size:
+            try:
+                content_range = response.info()["Content-Range"]
+                file_size = int(content_range.split("/")[1])
+            except (KeyError, IndexError, ValueError) as e:
+                logger.error(e)
         while True:
             chunk = response.read(chunk_size)
             if not chunk:

pytube/streams.py

@@ -9,9 +9,11 @@ has been renamed to accommodate DASH (which serves the audio and video
 separately).
 """
 
+from datetime import datetime
 import logging
 import os
 from typing import Dict, Tuple, Optional, BinaryIO
+from urllib.parse import parse_qs
 
 from pytube import extract
 from pytube import request
@@ -169,6 +171,11 @@ class Stream:
 
         return self.filesize
 
+    @property
+    def expiration(self) -> datetime:
+        expire = parse_qs(self.url.split("?")[1])["expire"][0]
+        return datetime.utcfromtimestamp(int(expire))
+
     @property
     def default_filename(self) -> str:
         """Generate filename based on the video title.

tests/test_request.py

@@ -7,9 +7,8 @@ import pytest
 from pytube import request
 
 
-@mock.patch("pytube.request.filesize", return_value=3 * 8 * 1024)
 @mock.patch("pytube.request.urlopen")
-def test_streaming(mock_urlopen, filesize):
+def test_streaming(mock_urlopen):
     # Given
     fake_stream_binary = [
         os.urandom(8 * 1024),
@@ -19,6 +18,7 @@ def test_streaming(mock_urlopen, filesize):
     ]
     response = mock.Mock()
     response.read.side_effect = fake_stream_binary
+    response.info.return_value = {"Content-Range": "bytes 200-1000/24576"}
     mock_urlopen.return_value = response
     # When
     response = request.stream("http://fakeassurl.gov")

tests/test_streams.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import os
 import random
+from datetime import datetime
 from unittest import mock
 from unittest.mock import MagicMock
 
@@ -59,6 +60,10 @@ def test_title(cipher_signature):
     assert cipher_signature.title == expected
 
 
+def test_expiration(cipher_signature):
+    assert cipher_signature.streams[0].expiration == datetime(2020, 1, 16, 5, 12, 5)
+
+
 def test_caption_tracks(presigned_video):
     assert len(presigned_video.caption_tracks) == 13