create Monostate class with callback protocols, add typing to __main__ and playlist

Makefile

@@ -10,8 +10,8 @@ pipenv:
 
 test:
 	pipenv run flake8
-	pipenv run black pytube --diff
-	pipenv run black tests --diff
+	pipenv run black pytube --check
+	pipenv run black tests --check
 	pipenv run mypy pytube
 	pipenv run pytest --cov-report term-missing --cov=pytube
 

pytube/__main__.py

@@ -10,6 +10,7 @@ smaller peripheral modules and functions.
 
 import json
 import logging
+from typing import Optional, Dict, List
 from urllib.parse import parse_qsl
 from html import unescape
 
@@ -23,6 +24,7 @@ from pytube import StreamQuery
 from pytube.mixins import install_proxy
 from pytube.exceptions import VideoUnavailable
 from pytube.helpers import apply_mixin
+from pytube.monostate import OnProgress, OnComplete, Monostate
 
 logger = logging.getLogger(__name__)
 
@@ -32,11 +34,11 @@ class YouTube(object):
 
     def __init__(
         self,
-        url,
-        defer_prefetch_init=False,
-        on_progress_callback=None,
-        on_complete_callback=None,
-        proxies=None,
+        url: str,
+        defer_prefetch_init: bool = False,
+        on_progress_callback: Optional[OnProgress] = None,
+        on_complete_callback: Optional[OnComplete] = None,
+        proxies: Dict[str, str] = None,
     ):
         """Construct a :class:`YouTube <YouTube>`.
 
@@ -52,23 +54,25 @@ class YouTube(object):
             complete events.
 
         """
-        self.js = None  # js fetched by js_url
-        self.js_url = None  # the url to the js, parsed from watch html
+        self.js: Optional[str] = None  # js fetched by js_url
+        self.js_url: Optional[str] = None  # the url to the js, parsed from watch html
 
         # note: vid_info may eventually be removed. It sounds like it once had
         # additional formats, but that doesn't appear to still be the case.
 
-        self.vid_info = None  # content fetched by vid_info_url
-        self.vid_info_url = None  # the url to vid info, parsed from watch html
+        # the url to vid info, parsed from watch html
+        self.vid_info_url: Optional[str] = None
+        self.vid_info_raw = None  # content fetched by vid_info_url
+        self.vid_info: Optional[Dict] = None  # parsed content of vid_info_raw
 
-        self.watch_html = None  # the html of /watch?v=<video_id>
-        self.embed_html = None
-        self.player_config_args = None  # inline js in the html containing
+        self.watch_html: Optional[str] = None  # the html of /watch?v=<video_id>
+        self.embed_html: Optional[str] = None
+        self.player_config_args: Dict = {}  # inline js in the html containing
         # streams
-        self.age_restricted = None
+        self.age_restricted: Optional[bool] = None
 
-        self.fmt_streams = []  # list of :class:`Stream <Stream>` instances
-        self.caption_tracks = []
+        self.fmt_streams: List[Stream] = []
+        self.caption_tracks: List[Caption] = []
 
         # video_id part of /watch?v=<video_id>
         self.video_id = extract.video_id(url)
@@ -79,11 +83,9 @@ class YouTube(object):
         self.embed_url = extract.embed_url(self.video_id)
         # A dictionary shared between all instances of :class:`Stream <Stream>`
         # (Borg pattern).
-        self.stream_monostate = {
-            # user defined callback functions.
-            "on_progress": on_progress_callback,
-            "on_complete": on_complete_callback,
-        }
+        self.stream_monostate = Monostate(
+            on_progress=on_progress_callback, on_complete=on_complete_callback
+        )
 
         if proxies:
             install_proxy(proxies)
@@ -113,16 +115,17 @@ class YouTube(object):
         """
         logger.info("init started")
 
-        self.vid_info = {k: v for k, v in parse_qsl(self.vid_info)}
+        self.vid_info = {k: v for k, v in parse_qsl(self.vid_info_raw)}
         if self.age_restricted:
             self.player_config_args = self.vid_info
         else:
+            assert self.watch_html is not None
             self.player_config_args = extract.get_ytplayer_config(self.watch_html,)[
                 "args"
             ]
 
             # Fix for KeyError: 'title' issue #434
-            if "title" not in self.player_config_args:
+            if "title" not in self.player_config_args:  # type: ignore
                 i_start = self.watch_html.lower().index("<title>") + len("<title>")
                 i_end = self.watch_html.lower().index("</title>")
                 title = self.watch_html[i_start:i_end].strip()
@@ -130,7 +133,8 @@ class YouTube(object):
                 title = title[:index] if index > 0 else title
                 self.player_config_args["title"] = unescape(title)
 
-        self.vid_descr = extract.get_vid_descr(self.watch_html)
+        if self.watch_html:
+            self.vid_descr = extract.get_vid_descr(self.watch_html)
         # https://github.com/nficano/pytube/issues/165
         stream_maps = ["url_encoded_fmt_stream_map"]
         if "adaptive_fmts" in self.player_config_args:
@@ -143,10 +147,14 @@ class YouTube(object):
             mixins.apply_descrambler(self.player_config_args, fmt)
 
             try:
-                mixins.apply_signature(self.player_config_args, fmt, self.js)
+                mixins.apply_signature(
+                    self.player_config_args, fmt, self.js  # type: ignore
+                )
             except TypeError:
-                self.js_url = extract.js_url(self.embed_html, self.age_restricted,)
+                assert self.embed_html is not None
+                self.js_url = extract.js_url(self.embed_html, self.age_restricted)
                 self.js = request.get(self.js_url)
+                assert self.js is not None
                 mixins.apply_signature(self.player_config_args, fmt, self.js)
 
             # build instances of :class:`Stream <Stream>`
@@ -169,18 +177,21 @@ class YouTube(object):
 
         """
         self.watch_html = request.get(url=self.watch_url)
-        if '<img class="icon meh" src="/yts/img' not in self.watch_html:
+        if (
+            self.watch_html is None
+            or '<img class="icon meh" src="/yts/img'  # noqa: W503
+            not in self.watch_html  # noqa: W503
+        ):
             raise VideoUnavailable("This video is unavailable.")
         self.embed_html = request.get(url=self.embed_url)
         self.age_restricted = extract.is_age_restricted(self.watch_html)
         self.vid_info_url = extract.video_info_url(
             video_id=self.video_id,
             watch_url=self.watch_url,
-            watch_html=self.watch_html,
             embed_html=self.embed_html,
             age_restricted=self.age_restricted,
         )
-        self.vid_info = request.get(self.vid_info_url)
+        self.vid_info_raw = request.get(self.vid_info_url)
         if not self.age_restricted:
             self.js_url = extract.js_url(self.watch_html, self.age_restricted)
             self.js = request.get(self.js_url)
@@ -302,9 +313,11 @@ class YouTube(object):
         :rtype: str
 
         """
-        return self.player_config_args["player_response"]["videoDetails"][
-            "lengthSeconds"
-        ]
+        return (
+            self.player_config_args.get("player_response", {})
+            .get("videoDetails", {})
+            .get("lengthSeconds")
+        )
 
     @property
     def views(self) -> str:
@@ -330,7 +343,7 @@ class YouTube(object):
             .get("author", "unknown")
         )
 
-    def register_on_progress_callback(self, func):
+    def register_on_progress_callback(self, func: OnProgress):
         """Register a download progress callback function post initialization.
 
         :param callable func:
@@ -340,9 +353,9 @@ class YouTube(object):
         :rtype: None
 
         """
-        self.stream_monostate["on_progress"] = func
+        self.stream_monostate.on_progress = func
 
-    def register_on_complete_callback(self, func):
+    def register_on_complete_callback(self, func: OnComplete):
         """Register a download complete callback function post initialization.
 
         :param callable func:
@@ -351,4 +364,4 @@ class YouTube(object):
         :rtype: None
 
         """
-        self.stream_monostate["on_complete"] = func
+        self.stream_monostate.on_complete = func

pytube/cli.py

@@ -136,18 +136,6 @@ def display_progress_bar(
 
 
 def on_progress(stream, chunk, file_handle, bytes_remaining):
-    """On download progress callback function.
-
-    :param object stream:
-        An instance of :class:`Stream <Stream>` being downloaded.
-    :param file_handle:
-        The file handle where the media is being written to.
-    :type file_handle:
-        :py:class:`io.BufferedWriter`
-    :param int bytes_remaining:
-        How many bytes have been downloaded.
-
-    """
     filesize = stream.filesize
     bytes_received = filesize - bytes_remaining
     display_progress_bar(bytes_received, filesize)

pytube/contrib/playlist.py

@@ -6,6 +6,7 @@ import json
 import logging
 import re
 from collections import OrderedDict
+from typing import List, Optional
 
 from pytube import request
 from pytube.__main__ import YouTube
@@ -18,12 +19,12 @@ class Playlist(object):
     playlist
     """
 
-    def __init__(self, url, suppress_exception=False):
+    def __init__(self, url: str, suppress_exception: bool = False):
         self.playlist_url = url
-        self.video_urls = []
+        self.video_urls: List[str] = []
         self.suppress_exception = suppress_exception
 
-    def construct_playlist_url(self):
+    def construct_playlist_url(self) -> str:
         """There are two kinds of playlist urls in YouTube. One that contains
         watch?v= in URL, another one contains the "playlist?list=" portion. It
         is preferable to work with the later one.
@@ -53,7 +54,7 @@ class Playlist(object):
             load_more_url = ""
         return load_more_url
 
-    def parse_links(self):
+    def parse_links(self) -> List[str]:
         """Parse the video links from the page source, extracts and
         returns the /watch?v= part from video link href
         It's an alternative for BeautifulSoup
@@ -115,8 +116,11 @@ class Playlist(object):
         return (str(i).zfill(digits) for i in range(start, stop, step))
 
     def download_all(
-        self, download_path=None, prefix_number=True, reverse_numbering=False,
-    ):
+        self,
+        download_path: Optional[str] = None,
+        prefix_number: bool = True,
+        reverse_numbering: bool = False,
+    ) -> None:
         """Download all the videos in the the playlist. Initially, download
         resolution is 720p (or highest available), later more option
         should be added to download resolution of choice
@@ -172,7 +176,7 @@ class Playlist(object):
                     dl_stream.download(download_path)
                 logger.debug("download complete")
 
-    def title(self):
+    def title(self) -> Optional[str]:
         """return playlist title (name)
         """
         try:
@@ -180,14 +184,19 @@ class Playlist(object):
             req = request.get(url)
             open_tag = "<title>"
             end_tag = "</title>"
-            match_result = re.compile(open_tag + "(.+?)" + end_tag)
-            match_result = match_result.search(req).group()
-            match_result = match_result.replace(open_tag, "")
-            match_result = match_result.replace(end_tag, "")
-            match_result = match_result.replace("- YouTube", "")
-            match_result = match_result.strip()
-
-            return match_result
+            pattern = re.compile(open_tag + "(.+?)" + end_tag)
+            match = pattern.search(req)
+
+            if match is None:
+                return None
+
+            return (
+                match.group()
+                .replace(open_tag, "")
+                .replace(end_tag, "")
+                .replace("- YouTube", "")
+                .strip()
+            )
         except Exception as e:
             logger.debug(e)
             return None

pytube/extract.py

@@ -96,11 +96,7 @@ def eurl(video_id: str) -> str:
 
 
 def video_info_url(
-    video_id: str,
-    watch_url: str,
-    watch_html: Optional[str],
-    embed_html: str,
-    age_restricted: bool,
+    video_id: str, watch_url: str, embed_html: Optional[str], age_restricted: bool,
 ) -> str:
     """Construct the video_info url.
 
@@ -108,8 +104,6 @@ def video_info_url(
         A YouTube video identifier.
     :param str watch_url:
         A YouTube watch url.
-    :param str watch_html:
-        (Unused) The html contents of the watch page.
     :param str embed_html:
         The html contents of the embed page (for age restricted videos).
     :param bool age_restricted:
@@ -139,7 +133,7 @@ def video_info_url(
     return "https://youtube.com/get_video_info?" + urlencode(params)
 
 
-def js_url(html: str, age_restricted: bool = False) -> str:
+def js_url(html: str, age_restricted: Optional[bool] = False) -> str:
     """Get the base JavaScript url.
 
     Construct the base JavaScript url, which contains the decipher
@@ -151,7 +145,7 @@ def js_url(html: str, age_restricted: bool = False) -> str:
         Is video age restricted.
 
     """
-    ytplayer_config = get_ytplayer_config(html, age_restricted)
+    ytplayer_config = get_ytplayer_config(html, age_restricted or False)
     base_js = ytplayer_config["assets"]["js"]
     return "https://youtube.com" + base_js
 

pytube/monostate.py

@@ -0,0 +1,59 @@
+import io
+from typing import Any, Optional
+from typing_extensions import Protocol
+
+
+class OnProgress(Protocol):
+    def __call__(
+        self,
+        stream: Any,
+        chunk: Any,
+        file_handler: io.BufferedWriter,
+        bytes_remaining: int,
+    ) -> None:
+        ...
+
+    """On download progress callback function.
+
+    :param stream:
+        An instance of :class:`Stream <Stream>` being downloaded.
+    :type stream:
+        :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.
+
+    """
+
+
+class OnComplete(Protocol):
+    def __call__(self, stream: Any, file_handler: io.BufferedWriter) -> 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:
+        The file handle where the media is being written to.
+    :type file_handler:
+        :py:class:`io.BufferedWriter`
+
+    :rtype: None
+
+    """
+
+
+class Monostate:
+    def __init__(
+        self, on_progress: Optional[OnProgress], on_complete: Optional[OnComplete]
+    ):
+        self.on_progress = on_progress
+        self.on_complete = on_complete

pytube/streams.py

@@ -18,7 +18,7 @@ from pytube import extract
 from pytube import request
 from pytube.helpers import safe_filename
 from pytube.itags import get_format_profile
-
+from pytube.monostate import Monostate
 
 logger = logging.getLogger(__name__)
 
@@ -26,7 +26,7 @@ logger = logging.getLogger(__name__)
 class Stream(object):
     """Container for stream manifest data."""
 
-    def __init__(self, stream: Dict, player_config_args: Dict, monostate: Dict):
+    def __init__(self, stream: Dict, player_config_args: Dict, monostate: Monostate):
         """Construct a :class:`Stream <Stream>`.
 
         :param dict stream:
@@ -303,7 +303,7 @@ class Stream(object):
                 indent=2,
             ),
         )
-        on_progress = self._monostate["on_progress"]
+        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)
@@ -320,7 +320,7 @@ class Stream(object):
 
         """
         logger.debug("download finished")
-        on_complete = self._monostate["on_complete"]
+        on_complete = self._monostate.on_complete
         if on_complete:
             logger.debug("calling on_complete callback %s", on_complete)
             on_complete(self, file_handle)

tests/test_extract.py

@@ -19,7 +19,6 @@ def test_info_url(cipher_signature):
     video_info_url = extract.video_info_url(
         video_id=cipher_signature.video_id,
         watch_url=cipher_signature.watch_url,
-        watch_html=cipher_signature.watch_html,
         embed_html="",
         age_restricted=False,
     )