Merge pull request #54 from hbmartin/live-and-complete-fixes

pytube/__main__.py

@@ -22,7 +22,7 @@ from pytube import Stream
 from pytube import StreamQuery
 from pytube.extract import apply_descrambler, apply_signature, get_ytplayer_config
 from pytube.helpers import install_proxy
-from pytube.exceptions import VideoUnavailable, LiveStreamError
+from pytube.exceptions import VideoUnavailable
 from pytube.monostate import OnProgress, OnComplete, Monostate
 
 logger = logging.getLogger(__name__)
@@ -164,11 +164,8 @@ class YouTube:
             raise VideoUnavailable(video_id=self.video_id)
         self.age_restricted = extract.is_age_restricted(self.watch_html)
 
-        if not self.age_restricted:
-            if "yt-badge-live" in self.watch_html:
-                raise LiveStreamError(self.video_id)
-            if "This video is private" in self.watch_html:
-                raise VideoUnavailable(video_id=self.video_id)
+        if not self.age_restricted and "This video is private" in self.watch_html:
+            raise VideoUnavailable(video_id=self.video_id)
 
         if self.age_restricted:
             if not self.embed_html:

pytube/extract.py

@@ -221,15 +221,16 @@ def apply_signature(config_args: Dict, fmt: str, js: str) -> None:
     """
     cipher = Cipher(js=js)
     stream_manifest = config_args[fmt]
-    live_stream = (
-        json.loads(config_args["player_response"])
-        .get("playabilityStatus", {},)
-        .get("liveStreamability")
-    )
+
     for i, stream in enumerate(stream_manifest):
         try:
             url: str = stream["url"]
         except KeyError:
+            live_stream = (
+                json.loads(config_args["player_response"])
+                .get("playabilityStatus", {},)
+                .get("liveStreamability")
+            )
             if live_stream:
                 raise LiveStreamError("UNKNOWN")
         # 403 Forbidden fix.

pytube/monostate.py

@@ -33,17 +33,16 @@ class OnProgress(Protocol):
 
 
 class OnComplete(Protocol):
-    def __call__(self, stream: Any, file_handler: io.BufferedWriter) -> None:
+    def __call__(self, stream: Any, file_path: Optional[str]) -> None:
         """On download complete handler function.
 
         :param stream:
             An instance of :class:`Stream <Stream>` being downloaded.
         :type stream:
             :py:class:`pytube.Stream`
-        :param file_handler:
+        :param file_path:
             The file handle where the media is being written to.
-        :type file_handler:
-            :py:class:`io.BufferedWriter`
+        :type file_path: str
 
         :rtype: None
         """

pytube/streams.py

@@ -245,7 +245,7 @@ class Stream:
                 bytes_remaining -= len(chunk)
                 # send to the on_progress callback.
                 self.on_progress(chunk, fh, bytes_remaining)
-        self.on_complete(fh)
+        self.on_complete(file_path)
         return file_path
 
     def stream_to_buffer(self) -> io.BytesIO:
@@ -264,7 +264,7 @@ class Stream:
             bytes_remaining -= len(chunk)
             # send to the on_progress callback.
             self.on_progress(chunk, buffer, bytes_remaining)
-        self.on_complete(buffer)
+        self.on_complete(None)
         return buffer
 
     def on_progress(self, chunk, file_handler, bytes_remaining):
@@ -300,13 +300,12 @@ class Stream:
             logger.debug("calling on_progress callback %s", on_progress)
             on_progress(self, chunk, file_handler, bytes_remaining)
 
-    def on_complete(self, file_handle):
+    def on_complete(self, file_path: Optional[str]):
         """On download complete handler function.
 
-        :param file_handle:
+        :param file_path:
             The file handle where the media is being written to.
-        :type file_handle:
-            :py:class:`io.BufferedWriter`
+        :type file_handle: str
 
         :rtype: None
 
@@ -315,7 +314,7 @@ class Stream:
         on_complete = self._monostate.on_complete
         if on_complete:
             logger.debug("calling on_complete callback %s", on_complete)
-            on_complete(self, file_handle)
+            on_complete(self, file_path)
 
     def __repr__(self) -> str:
         """Printable object representation.