more type checking

README.md

@@ -229,6 +229,8 @@ Finally, if you're filing a bug report, the cli contains a switch called ``--bui
 
 Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
 
+To run code checking before a PR use ``make test``
+
 #### Virtual environment
 
 Virtual environment is setup with [pipenv](https://pipenv-fork.readthedocs.io/en/latest/) and can be automatically activated with [direnv](https://direnv.net/docs/installation.html)

pytube/__main__.py

@@ -91,7 +91,7 @@ class YouTube(object):
         if not defer_prefetch_init:
             self.prefetch_descramble()
 
-    def prefetch_descramble(self):
+    def prefetch_descramble(self) -> None:
         """Download data, descramble it, and build Stream instances.
 
         :rtype: None
@@ -100,7 +100,7 @@ class YouTube(object):
         self.prefetch()
         self.descramble()
 
-    def descramble(self):
+    def descramble(self) -> None:
         """Descramble the stream data and build Stream instances.
 
         The initialization process takes advantage of Python's
@@ -158,7 +158,7 @@ class YouTube(object):
         self.initialize_caption_objects()
         logger.info("init finished successfully")
 
-    def prefetch(self):
+    def prefetch(self) -> None:
         """Eagerly download all necessary data.
 
         Eagerly executes all necessary network requests so all other
@@ -185,7 +185,7 @@ class YouTube(object):
             self.js_url = extract.js_url(self.watch_html, self.age_restricted)
             self.js = request.get(self.js_url)
 
-    def initialize_stream_objects(self, fmt: str):
+    def initialize_stream_objects(self, fmt: str) -> None:
         """Convert manifest data to instances of :class:`Stream <Stream>`.
 
         Take the unscrambled stream data and uses it to initialize
@@ -208,7 +208,7 @@ class YouTube(object):
             )
             self.fmt_streams.append(video)
 
-    def initialize_caption_objects(self):
+    def initialize_caption_objects(self) -> None:
         """Populate instances of :class:`Caption <Caption>`.
 
         Take the unscrambled player response data, and use it to initialize
@@ -230,7 +230,7 @@ class YouTube(object):
             self.caption_tracks.append(Caption(caption_track))
 
     @property
-    def captions(self):
+    def captions(self) -> CaptionQuery:
         """Interface to query caption tracks.
 
         :rtype: :class:`CaptionQuery <CaptionQuery>`.
@@ -238,7 +238,7 @@ class YouTube(object):
         return CaptionQuery([c for c in self.caption_tracks])
 
     @property
-    def streams(self):
+    def streams(self) -> StreamQuery:
         """Interface to query both adaptive (DASH) and progressive streams.
 
         :rtype: :class:`StreamQuery <StreamQuery>`.
@@ -246,7 +246,7 @@ class YouTube(object):
         return StreamQuery([s for s in self.fmt_streams])
 
     @property
-    def thumbnail_url(self):
+    def thumbnail_url(self) -> str:
         """Get the thumbnail url image.
 
         :rtype: str
@@ -255,7 +255,7 @@ class YouTube(object):
         return self.player_config_args["thumbnail_url"]
 
     @property
-    def title(self):
+    def title(self) -> str:
         """Get the video title.
 
         :rtype: str
@@ -264,7 +264,7 @@ class YouTube(object):
         return self.player_config_args["title"]
 
     @property
-    def description(self):
+    def description(self) -> str:
         """Get the video description.
 
         :rtype: str

pytube/captions.py

@@ -2,6 +2,7 @@
 import math
 import time
 import xml.etree.ElementTree as ElementTree
+from typing import Dict
 
 from pytube import request
 from html import unescape
@@ -10,7 +11,7 @@ from html import unescape
 class Caption:
     """Container for caption tracks."""
 
-    def __init__(self, caption_track):
+    def __init__(self, caption_track: Dict):
         """Construct a :class:`Caption <Caption>`.
 
         :param dict caption_track:
@@ -33,7 +34,7 @@ class Caption:
         """
         return self.xml_caption_to_srt(self.xml_captions)
 
-    def float_to_srt_time_format(self, d):
+    def float_to_srt_time_format(self, d: float) -> str:
         """Convert decimal durations into proper srt format.
 
         :rtype: str

pytube/cli.py

@@ -8,6 +8,7 @@ import json
 import logging
 import os
 import sys
+from typing import Tuple
 
 from pytube import __version__
 from pytube import YouTube
@@ -66,7 +67,7 @@ def main():
         download(args.url, args.itag)
 
 
-def build_playback_report(url: str):
+def build_playback_report(url: str) -> None:
     """Serialize the request data to json for offline debugging.
 
     :param str url:
@@ -95,13 +96,15 @@ def build_playback_report(url: str):
         )
 
 
-def get_terminal_size():
+def get_terminal_size() -> Tuple[int, int]:
     """Return the terminal size in rows and columns."""
     rows, columns = os.popen("stty size", "r").read().split()
     return int(rows), int(columns)
 
 
-def display_progress_bar(bytes_received, filesize, ch="█", scale=0.55):
+def display_progress_bar(
+    bytes_received: int, filesize: int, ch: str = "█", scale: float = 0.55
+) -> None:
     """Display a simple, pretty progress bar.
 
     Example:
@@ -150,7 +153,7 @@ def on_progress(stream, chunk, file_handle, bytes_remaining):
     display_progress_bar(bytes_received, filesize)
 
 
-def download(url: str, itag: str):
+def download(url: str, itag: str) -> None:
     """Start downloading a YouTube video.
 
     :param str url:
@@ -162,7 +165,10 @@ def download(url: str, itag: str):
     # TODO(nficano): allow download target to be specified
     # TODO(nficano): allow dash itags to be selected
     yt = YouTube(url, on_progress_callback=on_progress)
-    stream = yt.streams.get_by_itag(itag)
+    stream = yt.streams.get_by_itag(int(itag))
+    if stream is None:
+        print("Could not find a stream with itag: " + itag)
+        sys.exit()
     print("\n{fn} | {fs} bytes".format(fn=stream.default_filename, fs=stream.filesize,))
     try:
         stream.download()
@@ -171,7 +177,7 @@ def download(url: str, itag: str):
         sys.exit()
 
 
-def display_streams(url: str):
+def display_streams(url: str) -> None:
     """Probe YouTube video and lists its available formats.
 
     :param str url:

pytube/exceptions.py

@@ -15,7 +15,7 @@ class PytubeError(Exception):
 class ExtractError(PytubeError):
     """Data extraction based exception."""
 
-    def __init__(self, msg: str, video_id: str = None):
+    def __init__(self, msg: str, video_id: str = "unknown id"):
         """Construct an instance of a :class:`ExtractError <ExtractError>`.
 
         :param str msg:

pytube/extract.py

@@ -4,7 +4,7 @@ import json
 from collections import OrderedDict
 
 from html.parser import HTMLParser
-from typing import Any
+from typing import Any, Optional
 from urllib.parse import quote
 from urllib.parse import urlencode
 from pytube.exceptions import RegexMatchError
@@ -95,7 +95,7 @@ def eurl(video_id: str) -> str:
 def video_info_url(
     video_id: str,
     watch_url: str,
-    watch_html: str,
+    watch_html: Optional[str],
     embed_html: str,
     age_restricted: bool,
 ) -> str:

pytube/itags.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 """This module contains a lookup table of YouTube's itag values."""
+from typing import Dict
 
 ITAGS = {
     5: ("240p", "64kbps"),

pytube/mixins.py

@@ -17,7 +17,7 @@ from pytube.exceptions import LiveStreamError
 logger = logging.getLogger(__name__)
 
 
-def apply_signature(config_args: Dict, fmt: str, js: str):
+def apply_signature(config_args: Dict, fmt: str, js: str) -> None:
     """Apply the decrypted signature to the stream manifest.
 
     :param dict config_args:
@@ -67,7 +67,7 @@ def apply_signature(config_args: Dict, fmt: str, js: str):
         stream_manifest[i]["url"] = url + "&sig=" + signature
 
 
-def apply_descrambler(stream_data: Dict, key: str):
+def apply_descrambler(stream_data: Dict, key: str) -> None:
     """Apply various in-place transforms to YouTube's media stream data.
 
     Creates a ``list`` of dictionaries by string splitting on commas, then

pytube/query.py

@@ -206,7 +206,7 @@ class StreamQuery:
                 sorted(has_attribute, key=lambda s: getattr(s, attribute_name))
             )
 
-    def desc(self):
+    def desc(self) -> "StreamQuery":
         """Sort streams in descending order.
 
         :rtype: :class:`StreamQuery <StreamQuery>`
@@ -214,7 +214,7 @@ class StreamQuery:
         """
         return StreamQuery(self.fmt_streams[::-1])
 
-    def asc(self):
+    def asc(self) -> "StreamQuery":
         """Sort streams in ascending order.
 
         :rtype: :class:`StreamQuery <StreamQuery>`
@@ -222,7 +222,7 @@ class StreamQuery:
         """
         return self
 
-    def get_by_itag(self, itag):
+    def get_by_itag(self, itag) -> Optional[Stream]:
         """Get the corresponding :class:`Stream <Stream>` for a given itag.
 
         :param int itag:
@@ -233,12 +233,9 @@ class StreamQuery:
             not found.
 
         """
-        try:
-            return self.itag_index[int(itag)]
-        except KeyError:
-            pass
+        return self.itag_index.get(int(itag))
 
-    def first(self):
+    def first(self) -> Optional[Stream]:
         """Get the first :class:`Stream <Stream>` in the results.
 
         :rtype: :class:`Stream <Stream>` or None
@@ -250,7 +247,7 @@ class StreamQuery:
         try:
             return self.fmt_streams[0]
         except IndexError:
-            pass
+            return None
 
     def last(self):
         """Get the last :class:`Stream <Stream>` in the results.

pytube/streams.py

@@ -83,7 +83,7 @@ class Stream(object):
         # streams return NoneType for audio/video depending.
         self.video_codec, self.audio_codec = self.parse_codecs()
 
-    def set_attributes_from_dict(self, dct: Dict):
+    def set_attributes_from_dict(self, dct: Dict) -> None:
         """Set class attributes from dictionary items.
 
         :rtype: None
@@ -199,7 +199,12 @@ class Stream(object):
         filename = safe_filename(self.title)
         return "{filename}.{s.subtype}".format(filename=filename, s=self)
 
-    def download(self, output_path=None, filename=None, filename_prefix=None):
+    def download(
+        self,
+        output_path: Optional[str] = None,
+        filename: Optional[str] = None,
+        filename_prefix: Optional[str] = None,
+    ) -> str:
         """Write the media stream to disk.
 
         :param output_path:
@@ -248,7 +253,7 @@ class Stream(object):
         self.on_complete(fh)
         return fp
 
-    def stream_to_buffer(self):
+    def stream_to_buffer(self) -> io.BytesIO:
         """Write the media stream to buffer
 
         :rtype: io.BytesIO buffer