type hinting for query, streams, and extract

README.md

@@ -2,7 +2,7 @@
 <div align="center">
   <p align="center">
 	  <img src="https://img.shields.io/pypi/v/pytube3.svg" alt="pypi">
-	  <a href="https://travis-ci.org/hbmartin/pytube3"><img src="https://travis-ci.org/hbmartin/pytube3.svg?branch=master" /></a>
+	  <a href="https://travis-ci.com/hbmartin/pytube3/"><img src="https://travis-ci.org/hbmartin/pytube3.svg?branch=master" /></a>
 	  <a href='https://pytube3.readthedocs.io/en/latest/?badge=latest'><img src='https://readthedocs.org/projects/pytube3/badge/?version=latest' alt='Documentation Status' /></a>
 	  <a href="https://coveralls.io/github/hbmartin/pytube3?branch=master"><img src="https://coveralls.io/repos/github/hbmartin/pytube3/badge.svg?branch=master" /></a>
 	  <a href="https://pypi.python.org/pypi/pytube3/"><img src="https://img.shields.io/pypi/pyversions/pytube3.svg" /></a>
@@ -239,11 +239,11 @@ Virtual environment is setup with [pipenv](https://pipenv-fork.readthedocs.io/en
 
 #### Code Formatting
 
-This project is linted with [pyflakes](https://github.com/PyCQA/pyflakes) and makes strict use of [Black](https://github.com/ambv/black) for code formatting.
+This project is linted with [pyflakes](https://github.com/PyCQA/pyflakes), formatted with [Black](https://github.com/ambv/black), and typed with [mypy](https://mypy.readthedocs.io/en/latest/introduction.html)
 
 #### Testing
 
-This project is tested with [pytest](https://docs.pytest.org/en/latest/) and coverage is evaluated with [coveralls](https://coveralls-python.readthedocs.io/en/latest/index.html)
+This project is tested with [pytest](https://docs.pytest.org/en/latest/) and coverage is evaluated with [coveralls](https://coveralls-python.readthedocs.io/en/latest/index.html).
 
 #### Code of Conduct
 

pytube/extract.py

@@ -4,6 +4,7 @@ import json
 from collections import OrderedDict
 
 from html.parser import HTMLParser
+from typing import Any
 from urllib.parse import quote
 from urllib.parse import urlencode
 from pytube.exceptions import RegexMatchError
@@ -177,14 +178,14 @@ def mime_type_codec(mime_type_codec):
     return mime_type, [c.strip() for c in codecs.split(",")]
 
 
-def get_ytplayer_config(html, age_restricted=False):
+def get_ytplayer_config(html: str, age_restricted: bool = False) -> Any:
     """Get the YouTube player configuration data from the watch html.
 
     Extract the ``ytplayer_config``, which is json data embedded within the
     watch html and serves as the primary source of obtaining the stream
     manifest data.
 
-    :param str watch_html:
+    :param str html:
         The html contents of the watch page.
     :param bool age_restricted:
         Is video age restricted.
@@ -200,7 +201,7 @@ def get_ytplayer_config(html, age_restricted=False):
     return json.loads(yt_player_config)
 
 
-def get_vid_descr(html):
+def get_vid_descr(html: str) -> str:
     html_parser = PytubeHTMLParser()
     html_parser.feed(html)
     return html_parser.vid_descr

pytube/query.py

@@ -1,5 +1,8 @@
 # -*- coding: utf-8 -*-
 """This module provides a query interface for media streams and captions."""
+from typing import List, Optional
+
+from pytube import Stream, Caption
 
 
 class StreamQuery:
@@ -210,7 +213,7 @@ class StreamQuery:
     def get_by_itag(self, itag):
         """Get the corresponding :class:`Stream <Stream>` for a given itag.
 
-        :param str int itag:
+        :param int itag:
             YouTube format identifier code.
         :rtype: :class:`Stream <Stream>` or None
         :returns:
@@ -251,7 +254,7 @@ class StreamQuery:
         except IndexError:
             pass
 
-    def count(self):
+    def count(self) -> int:
         """Get the count the query would return.
 
         :rtype: int
@@ -259,7 +262,7 @@ class StreamQuery:
         """
         return len(self.fmt_streams)
 
-    def all(self):
+    def all(self) -> List[Stream]:
         """Get all the results represented by this query as a list.
 
         :rtype: list
@@ -271,7 +274,7 @@ class StreamQuery:
 class CaptionQuery:
     """Interface for querying the available captions."""
 
-    def __init__(self, captions):
+    def __init__(self, captions: List[Caption]):
         """Construct a :class:`Caption <Caption>`.
 
         param list captions:
@@ -281,7 +284,7 @@ class CaptionQuery:
         self.captions = captions
         self.lang_code_index = {c.code: c for c in captions}
 
-    def get_by_language_code(self, lang_code):
+    def get_by_language_code(self, lang_code: str) -> Optional[Caption]:
         """Get the :class:`Caption <Caption>` for a given ``lang_code``.
 
         :param str lang_code:
@@ -293,7 +296,7 @@ class CaptionQuery:
         """
         return self.lang_code_index.get(lang_code)
 
-    def all(self):
+    def all(self) -> List[Caption]:
         """Get all the results represented by this query as a list.
 
         :rtype: list

pytube/streams.py

@@ -12,6 +12,7 @@ import io
 import logging
 import os
 import pprint
+from typing import Dict, Tuple, Optional
 
 from pytube import extract
 from pytube import request
@@ -25,7 +26,7 @@ logger = logging.getLogger(__name__)
 class Stream(object):
     """Container for stream manifest data."""
 
-    def __init__(self, stream, player_config_args, monostate):
+    def __init__(self, stream: Dict, player_config_args: Dict, monostate: Dict):
         """Construct a :class:`Stream <Stream>`.
 
         :param dict stream:
@@ -47,7 +48,7 @@ class Stream(object):
         self.res = None  # resolution (e.g.: 480p, 720p, 1080p)
         self.url = None  # signed download url
 
-        self._filesize = None  # filesize in bytes
+        self._filesize: Optional[int] = None  # filesize in bytes
         self.mime_type = None  # content identifier (e.g.: video/mp4)
         self.type = None  # the part of the mime before the slash
         self.subtype = None  # the part of the mime after the slash
@@ -82,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):
+    def set_attributes_from_dict(self, dct: Dict):
         """Set class attributes from dictionary items.
 
         :rtype: None
@@ -91,17 +92,17 @@ class Stream(object):
             setattr(self, key, val)
 
     @property
-    def is_adaptive(self):
+    def is_adaptive(self) -> bool:
         """Whether the stream is DASH.
 
         :rtype: bool
         """
         # if codecs has two elements (e.g.: ['vp8', 'vorbis']): 2 % 2 = 0
         # if codecs has one element (e.g.: ['vp8']) 1 % 2 = 1
-        return len(self.codecs) % 2
+        return bool(len(self.codecs) % 2)
 
     @property
-    def is_progressive(self):
+    def is_progressive(self) -> bool:
         """Whether the stream is progressive.
 
         :rtype: bool
@@ -109,7 +110,7 @@ class Stream(object):
         return not self.is_adaptive
 
     @property
-    def includes_audio_track(self):
+    def includes_audio_track(self) -> bool:
         """Whether the stream only contains audio.
 
         :rtype: bool
@@ -119,7 +120,7 @@ class Stream(object):
         return self.type == "audio"
 
     @property
-    def includes_video_track(self):
+    def includes_video_track(self) -> bool:
         """Whether the stream only contains video.
 
         :rtype: bool
@@ -128,7 +129,7 @@ class Stream(object):
             return True
         return self.type == "video"
 
-    def parse_codecs(self):
+    def parse_codecs(self) -> Tuple:
         """Get the video/audio codecs from list of codecs.
 
         Parse a variable length sized list of codecs and returns a
@@ -152,7 +153,7 @@ class Stream(object):
         return video, audio
 
     @property
-    def filesize(self):
+    def filesize(self) -> int:
         """File size of the media stream in bytes.
 
         :rtype: int
@@ -165,7 +166,7 @@ class Stream(object):
         return self._filesize
 
     @property
-    def title(self):
+    def title(self) -> str:
         """Get title of video
 
         :rtype: str
@@ -187,7 +188,7 @@ class Stream(object):
         return "Unknown YouTube Video Title"
 
     @property
-    def default_filename(self):
+    def default_filename(self) -> str:
         """Generate filename based on the video title.
 
         :rtype: str
@@ -316,7 +317,7 @@ class Stream(object):
             logger.debug("calling on_complete callback %s", on_complete)
             on_complete(self, file_handle)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Printable object representation.
 
         :rtype: str
@@ -335,5 +336,4 @@ class Stream(object):
                 parts.extend(['vcodec="{s.video_codec}"'])
         else:
             parts.extend(['abr="{s.abr}"', 'acodec="{s.audio_codec}"'])
-        parts = " ".join(parts).format(s=self)
-        return "<Stream: {parts}>".format(parts=parts)
+        return "<Stream: {parts}>".format(parts=" ".join(parts).format(s=self))