multi range file streaming

README.md

@@ -36,7 +36,7 @@ $ pip install pytube3 --upgrade
 ## Quick start
 ```python
  >>> from pytube import YouTube
- >>> YouTube('https://youtu.be/9bZkp7q19f0').streams[0].download()
+ >>> YouTube('https://youtu.be/9bZkp7q19f0').streams.get_highest_resolution().download()
  >>>
  >>> yt = YouTube('http://youtube.com/watch?v=9bZkp7q19f0')
  >>> yt.streams

pytube/cli.py

@@ -11,7 +11,6 @@ import os
 import shutil
 import sys
 import subprocess  # nosec
-from io import BufferedWriter
 from typing import Any, Optional, List
 
 from pytube import __version__, CaptionQuery, Stream, Playlist
@@ -210,9 +209,7 @@ def display_progress_bar(
 
 
 # noinspection PyUnusedLocal
-def on_progress(
-    stream: Any, chunk: Any, file_handler: BufferedWriter, bytes_remaining: int
-) -> None:
+def on_progress(stream: Any, chunk: bytes, bytes_remaining: int) -> None:
     filesize = stream.filesize
     bytes_received = filesize - bytes_remaining
     display_progress_bar(bytes_received, filesize)

pytube/helpers.py

@@ -4,7 +4,6 @@
 import functools
 import logging
 import os
-import pprint
 import re
 import warnings
 from typing import TypeVar, Callable, Optional, Dict, List, Any
@@ -34,10 +33,7 @@ def regex_search(pattern: str, string: str, group: int) -> str:
     if not results:
         raise RegexMatchError(caller="regex_search", pattern=pattern)
 
-    logger.debug(
-        "finished regex search: %s",
-        pprint.pformat({"pattern": pattern, "results": results.group(0),}, indent=2,),
-    )
+    logger.debug("matched regex search: %s", pattern)
 
     return results.group(group)
 

pytube/monostate.py

@@ -1,18 +1,11 @@
 # -*- coding: utf-8 -*-
 
-import io
 from typing import Any, Optional
 from typing_extensions import Protocol
 
 
 class OnProgress(Protocol):
-    def __call__(
-        self,
-        stream: bytes,
-        chunk: Any,
-        file_handler: io.BufferedWriter,
-        bytes_remaining: int,
-    ) -> None:
+    def __call__(self, stream: Any, chunk: bytes, bytes_remaining: int) -> None:
         """On download progress callback function.
 
         :param stream:
@@ -21,10 +14,6 @@ class OnProgress(Protocol):
             :py:class:`pytube.Stream`
         :param str chunk:
             Segment of media file binary data, not yet written to disk.
-        :param file_handler:
-            The file handle where the media is being written to.
-        :type file_handler:
-            :py:class:`io.BufferedWriter`
         :param int bytes_remaining:
             How many bytes have been downloaded.
 

pytube/request.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 """Implements a simple wrapper around urlopen."""
+from functools import lru_cache
 from typing import Any, Iterable, Dict, Optional
 from urllib.request import Request
 from urllib.request import urlopen
@@ -31,18 +32,33 @@ def get(url) -> str:
     return _execute_request(url).read().decode("utf-8")
 
 
-def stream(url: str, chunk_size: int = 8192) -> Iterable[bytes]:
+def stream(
+    url: str, chunk_size: int = 1024, range_size: int = 10485760
+) -> 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
+    :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
     :rtype: Iterable[bytes]
     """
-    response = _execute_request(url, headers={"Range": "bytes=0-"})
-    while True:
-        buf = response.read(chunk_size)
-        if not buf:
-            break
-        yield buf
+    file_size: int = filesize(url)
+    downloaded = 0
+    while downloaded < file_size:
+        stop_pos = min(downloaded + range_size, file_size) - 1
+        range_header = f"bytes={downloaded}-{stop_pos}"
+        r = _execute_request(url, method="GET", headers={"Range": range_header})
+        while True:
+            chunk = r.read(chunk_size)
+            if not chunk:
+                break
+            downloaded += len(chunk)
+            yield chunk
+    return
+
+
+@lru_cache(maxsize=None)
+def filesize(url: str) -> int:
+    return int(head(url)["content-length"])
 
 
 def head(url: str) -> Dict:

pytube/streams.py

@@ -9,11 +9,9 @@ has been renamed to accommodate DASH (which serves the audio and video
 separately).
 """
 
-import io
 import logging
 import os
-import pprint
-from typing import Dict, Tuple, Optional
+from typing import Dict, Tuple, Optional, BinaryIO
 
 from pytube import extract
 from pytube import request
@@ -143,8 +141,7 @@ class Stream:
             Filesize (in bytes) of the stream.
         """
         if self._filesize is None:
-            headers = request.head(self.url)
-            self._filesize = int(headers["content-length"])
+            self._filesize = request.filesize(self.url)
         return self._filesize
 
     @property
@@ -255,15 +252,14 @@ class Stream:
     def exists_at_path(self, file_path: str) -> bool:
         return os.path.isfile(file_path) and os.path.getsize(file_path) == self.filesize
 
-    def stream_to_buffer(self) -> io.BytesIO:
+    def stream_to_buffer(self, buffer: BinaryIO) -> None:
         """Write the media stream to buffer
 
         :rtype: io.BytesIO buffer
         """
-        buffer = io.BytesIO()
         bytes_remaining = self.filesize
-        logger.debug(
-            "downloading (%s total bytes) file to BytesIO buffer", self.filesize,
+        logger.info(
+            "downloading (%s total bytes) file to buffer", self.filesize,
         )
 
         for chunk in request.stream(self.url):
@@ -272,9 +268,8 @@ class Stream:
             # send to the on_progress callback.
             self.on_progress(chunk, buffer, bytes_remaining)
         self.on_complete(None)
-        return buffer
 
-    def on_progress(self, chunk, file_handler, bytes_remaining):
+    def on_progress(self, chunk: bytes, file_handler: BinaryIO, bytes_remaining: int):
         """On progress callback function.
 
         This function writes the binary data to the file, then checks if an
@@ -295,24 +290,16 @@ class Stream:
 
         """
         file_handler.write(chunk)
-        logger.debug(
-            "download progress\n%s",
-            pprint.pformat(
-                {"chunk_size": len(chunk), "bytes_remaining": bytes_remaining,},
-                indent=2,
-            ),
-        )
-        on_progress = self._monostate.on_progress
-        if on_progress:
-            logger.debug("calling on_progress callback %s", on_progress)
-            on_progress(self, chunk, file_handler, bytes_remaining)
+        logger.debug("download remaining: %s", bytes_remaining)
+        if self._monostate.on_progress:
+            self._monostate.on_progress(self, chunk, bytes_remaining)
 
     def on_complete(self, file_path: Optional[str]):
         """On download complete handler function.
 
         :param file_path:
             The file handle where the media is being written to.
-        :type file_handle: str
+        :type file_path: str
 
         :rtype: None
 

tests/test_cli.py

@@ -125,11 +125,10 @@ def test_display_progress_bar(capsys):
 
 
 @mock.patch("pytube.Stream")
-@mock.patch("io.BufferedWriter")
-def test_on_progress(stream, writer):
+def test_on_progress(stream):
     stream.filesize = 10
     cli.display_progress_bar = MagicMock()
-    cli.on_progress(stream, "", writer, 7)
+    cli.on_progress(stream, "", 7)
     cli.display_progress_bar.assert_called_once_with(3, 10)
 
 

tests/test_request.py

@@ -7,12 +7,13 @@ 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):
+def test_streaming(mock_urlopen, filesize):
     fake_stream_binary = [
-        iter(os.urandom(8 * 1024)),
-        iter(os.urandom(8 * 1024)),
-        iter(os.urandom(8 * 1024)),
+        os.urandom(8 * 1024),
+        os.urandom(8 * 1024),
+        os.urandom(8 * 1024),
         None,
     ]
     response = mock.Mock()

tests/test_streams.py

@@ -15,12 +15,11 @@ def test_filesize(cipher_signature, mocker):
 
 
 def test_filesize_approx(cipher_signature, mocker):
-    mocker.patch.object(request, "head")
-    request.head.return_value = {"content-length": "123"}
     stream = cipher_signature.streams[0]
+
     assert stream.filesize_approx == 22350604
     stream.bitrate = None
-    assert stream.filesize_approx == 123
+    assert stream.filesize_approx == 6796391
 
 
 def test_default_filename(cipher_signature):
@@ -188,8 +187,8 @@ def test_on_progress_hook(cipher_signature, mocker):
         stream.download()
     assert callback_fn.called
     args, _ = callback_fn.call_args
-    assert len(args) == 4
-    stream, _, _, _ = args
+    assert len(args) == 3
+    stream, _, _ = args
     assert isinstance(stream, Stream)