add -r resolution download tag

pytube/cli.py

@@ -11,7 +11,7 @@ import sys
 from io import BufferedWriter
 from typing import Tuple, Any, Optional, List
 
-from pytube import __version__, CaptionQuery
+from pytube import __version__, CaptionQuery, Stream
 from pytube import YouTube
 
 
@@ -36,9 +36,11 @@ def main():
     if args.build_playback_report:
         build_playback_report(youtube)
     if args.itag:
-        download(youtube=youtube, itag=args.itag)
+        download_by_itag(youtube=youtube, itag=args.itag)
     if hasattr(args, "caption_code"):
         download_caption(youtube=youtube, lang_code=args.caption_code)
+    if args.resolution:
+        download_by_resolution(youtube=youtube, resolution=args.resolution)
 
 
 def _parse_args(
@@ -51,6 +53,9 @@ def _parse_args(
     parser.add_argument(
         "--itag", type=int, help="The itag for the desired stream",
     )
+    parser.add_argument(
+        "-r", "--resolution", type=str, help="The resolution for the desired stream",
+    )
     parser.add_argument(
         "-l",
         "--list",
@@ -165,12 +170,18 @@ def on_progress(
     display_progress_bar(bytes_received, filesize)
 
 
-def download(youtube: YouTube, itag: int) -> None:
+def _download(stream: Stream) -> None:
+    print("\n{fn} | {fs} bytes".format(fn=stream.default_filename, fs=stream.filesize, ))
+    stream.download()
+    sys.stdout.write("\n")
+
+
+def download_by_itag(youtube: YouTube, itag: int) -> None:
     """Start downloading a YouTube video.
 
     :param YouTube youtube:
         A valid YouTube object.
-    :param str itag:
+    :param int itag:
         YouTube format identifier code.
 
     """
@@ -184,10 +195,35 @@ def download(youtube: YouTube, itag: int) -> None:
         sys.exit()
 
     youtube.register_on_progress_callback(on_progress)
-    print("\n{fn} | {fs} bytes".format(fn=stream.default_filename, fs=stream.filesize,))
+
+    try:
+        _download(stream)
+    except KeyboardInterrupt:
+        sys.exit()
+
+
+def download_by_resolution(youtube: YouTube, resolution: str) -> None:
+    """Start downloading a YouTube video.
+
+    :param YouTube youtube:
+        A valid YouTube object.
+    :param str resolution:
+        YouTube video resolution.
+
+    """
+    # TODO(nficano): allow download target to be specified
+    # TODO(nficano): allow dash itags to be selected
+    stream = youtube.streams.get_by_resolution(resolution)
+    if stream is None:
+        print("Could not find a stream with resolution: {resolution}".format(resolution=resolution))
+        print("Try one of these:")
+        display_streams(youtube)
+        sys.exit()
+
+    youtube.register_on_progress_callback(on_progress)
+
     try:
-        stream.download()
-        sys.stdout.write("\n")
+        _download(stream)
     except KeyboardInterrupt:
         sys.exit()
 

pytube/contrib/playlist.py

@@ -19,25 +19,14 @@ class Playlist:
     """
 
     def __init__(self, url: str, suppress_exception: bool = False):
-        self.playlist_url = url
         self.video_urls: List[str] = []
         self.suppress_exception = suppress_exception
+        self.playlist_url:str = url
 
-    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.
-
-        :return: playlist url
-        """
-
-        if "watch?v=" in self.playlist_url:
+        if "watch?v=" in url:
             base_url = "https://www.youtube.com/playlist?list="
-            playlist_code = self.playlist_url.split("&list=")[1]
-            return base_url + playlist_code
-
-        # url is already in the desired format, so just return it
-        return self.playlist_url
+            playlist_code = self.playlist_url.split("&list=")[1]  # TODO: should be parse q
+            self.playlist_url = base_url + playlist_code
 
     @staticmethod
     def _find_load_more_url(req: str) -> Optional[str]:
@@ -59,8 +48,7 @@ class Playlist:
         It's an alternative for BeautifulSoup
         """
 
-        url = self.construct_playlist_url()
-        req = request.get(url)
+        req = request.get(self.playlist_url)
 
         # split the page source by line and process each line
         content = [x for x in req.split("\n") if "pl-video-title-link" in x]
@@ -176,8 +164,7 @@ class Playlist:
 
     def title(self) -> Optional[str]:
         """return playlist title (name)"""
-        url = self.construct_playlist_url()
-        req = request.get(url)
+        req = request.get(self.playlist_url)
         open_tag = "<title>"
         end_tag = "</title>"
         pattern = re.compile(open_tag + "(.+?)" + end_tag)

tests/test_cli.py

@@ -16,7 +16,7 @@ def test_download_when_itag_not_found(youtube):
     youtube.streams.all.return_value = []
     youtube.streams.get_by_itag.return_value = None
     with pytest.raises(SystemExit):
-        cli.download(youtube, 123)
+        cli.download_by_itag(youtube, 123)
     youtube.streams.get_by_itag.assert_called_with(123)
 
 
@@ -28,7 +28,7 @@ def test_download_when_itag_is_found(youtube, stream):
     with patch.object(
         youtube.streams, "get_by_itag", wraps=youtube.streams.get_by_itag
     ) as wrapped_itag:
-        cli.download(youtube, 123)
+        cli.download_by_itag(youtube, 123)
         wrapped_itag.assert_called_with(123)
     youtube.register_on_progress_callback.assert_called_with(cli.on_progress)
     stream.download.assert_called()
@@ -115,14 +115,14 @@ def test_parse_args_truthy():
 
 
 @mock.patch("pytube.cli.YouTube.__init__", return_value=None)
-def test_main_download(youtube):
+def test_main_download_by_itag(youtube):
     parser = argparse.ArgumentParser()
     args = parse_args(parser, ["urlhere", "--itag=10"])
     cli._parse_args = MagicMock(return_value=args)
-    cli.download = MagicMock()
+    cli.download_by_itag = MagicMock()
     cli.main()
     youtube.assert_called()
-    cli.download.assert_called()
+    cli.download_by_itag.assert_called()
 
 
 @mock.patch("pytube.cli.YouTube.__init__", return_value=None)