Merge pull request #7 from hbmartin/cli-enhancements

.deepsource.toml

@@ -1,7 +1,7 @@
 version = 1
 
 test_patterns = [
-  "*/tests/**"
+  "tests/**"
 ]
 
 exclude_patterns = [

README.md

@@ -232,7 +232,7 @@ pytube also ships with a tiny cli interface for downloading and probing videos.
 Let's start with downloading:
 
 ```bash
-$ pytube3 http://youtube.com/watch?v=9bZkp7q19f0 --itag=22
+$ pytube3 http://youtube.com/watch?v=9bZkp7q19f0 --itag=18
 ```
 To view available streams:
 

pytube/captions.py

@@ -1,12 +1,15 @@
 # -*- coding: utf-8 -*-
 import math
+import os
 import time
 import xml.etree.ElementTree as ElementTree
-from typing import Dict
+from typing import Dict, Optional
 
 from pytube import request
 from html import unescape
 
+from pytube.helpers import safe_filename
+
 
 class Caption:
     """Container for caption tracks."""
@@ -73,6 +76,68 @@ class Caption:
             segments.append(line)
         return "\n".join(segments).strip()
 
+    def download(
+        self,
+        title: str,
+        srt: bool = True,
+        output_path: Optional[str] = None,
+        filename_prefix: Optional[str] = None,
+    ) -> str:
+        """Write the media stream to disk.
+
+        :param filename:
+            Output filename (stem only) for writing media file.
+            If one is not specified, the default filename is used.
+        :type filename: str
+        :param srt:
+            Set to True to download srt, false to download xml. Defaults to True.
+        :type srt bool
+        :param output_path:
+            (optional) Output path for writing media file. If one is not
+            specified, defaults to the current working directory.
+        :type output_path: str or None
+        :param filename_prefix:
+            (optional) A string that will be prepended to the filename.
+            For example a number in a playlist or the name of a series.
+            If one is not specified, nothing will be prepended
+            This is separate from filename so you can use the default
+            filename but still add a prefix.
+        :type filename_prefix: str or None
+
+        :rtype: str
+
+        """
+        output_path = output_path or os.getcwd()
+
+        if title.endswith(".srt") or title.endswith(".xml"):
+            filename = ".".join(title.split(".")[:-1])
+        else:
+            filename = title
+
+        if filename_prefix:
+            filename = "{prefix}{filename}".format(
+                prefix=safe_filename(filename_prefix), filename=filename,
+            )
+
+        filename = safe_filename(filename)
+
+        filename += " ({})".format(self.code)
+
+        if srt:
+            filename += ".srt"
+        else:
+            filename += ".xml"
+
+        file_path = os.path.join(output_path, filename)
+
+        with open(file_path, "w", encoding="utf-8") as file_handle:
+            if srt:
+                file_handle.write(self.generate_srt_captions())
+            else:
+                file_handle.write(self.xml_captions)
+
+        return file_path
+
     def __repr__(self):
         """Printable object representation."""
         return '<Caption lang="{s.name}" code="{s.code}">'.format(s=self)

pytube/cli.py

@@ -1,4 +1,4 @@
-# -*- coding: utf-8 -*-
+#!/usr/bin/env python3
 """A simple command line application to download youtube videos."""
 
 import argparse
@@ -9,9 +9,9 @@ import logging
 import os
 import sys
 from io import BufferedWriter
-from typing import Tuple, Any
+from typing import Tuple, Any, Optional, List
 
-from pytube import __version__
+from pytube import __version__, CaptionQuery
 from pytube import YouTube
 
 
@@ -22,6 +22,28 @@ def main():
     """Command line application to download youtube videos."""
     # noinspection PyTypeChecker
     parser = argparse.ArgumentParser(description=main.__doc__)
+    args = _parse_args(parser)
+    logging.getLogger().setLevel(max(3 - args.verbosity, 0) * 10)
+
+    if not args.url:
+        parser.print_help()
+        sys.exit(1)
+
+    youtube = YouTube(args.url)
+
+    if args.list:
+        display_streams(youtube)
+    if args.build_playback_report:
+        build_playback_report(youtube)
+    if args.itag:
+        download(youtube=youtube, itag=args.itag)
+    if hasattr(args, "caption_code"):
+        download_caption(youtube=youtube, lang_code=args.caption_code)
+
+
+def _parse_args(
+    parser: argparse.ArgumentParser, args: Optional[List] = None
+) -> argparse.Namespace:
     parser.add_argument("url", help="The YouTube /watch url", nargs="?")
     parser.add_argument(
         "--version", action="version", version="%(prog)s " + __version__,
@@ -51,45 +73,41 @@ def main():
         action="store_true",
         help="Save the html and js to disk",
     )
+    parser.add_argument(
+        "-c",
+        "--caption-code",
+        type=str,
+        default=argparse.SUPPRESS,
+        nargs="?",
+        help=(
+            "Download srt captions for given language code. "
+            "Prints available language codes if no argument given"
+        ),
+    )
 
-    args = parser.parse_args()
-    logging.getLogger().setLevel(max(3 - args.verbosity, 0) * 10)
-
-    if not args.url:
-        parser.print_help()
-        sys.exit(1)
-
-    if args.list:
-        display_streams(args.url)
-
-    elif args.build_playback_report:
-        build_playback_report(args.url)
-
-    elif args.itag:
-        download(args.url, args.itag)
+    return parser.parse_args(args)
 
 
-def build_playback_report(url: str) -> None:
+def build_playback_report(youtube: YouTube) -> None:
     """Serialize the request data to json for offline debugging.
 
-    :param str url:
-        A valid YouTube watch URL.
+    :param YouTube youtube:
+        A YouTube object.
     """
-    yt = YouTube(url)
     ts = int(dt.datetime.utcnow().timestamp())
     fp = os.path.join(
-        os.getcwd(), "yt-video-{yt.video_id}-{ts}.json.gz".format(yt=yt, ts=ts),
+        os.getcwd(), "yt-video-{yt.video_id}-{ts}.json.gz".format(yt=youtube, ts=ts),
     )
 
-    js = yt.js
-    watch_html = yt.watch_html
-    vid_info = yt.vid_info
+    js = youtube.js
+    watch_html = youtube.watch_html
+    vid_info = youtube.vid_info
 
     with gzip.open(fp, "wb") as fh:
         fh.write(
             json.dumps(
                 {
-                    "url": url,
+                    "url": youtube.watch_url,
                     "js": js,
                     "watch_html": watch_html,
                     "video_info": vid_info,
@@ -147,22 +165,25 @@ def on_progress(
     display_progress_bar(bytes_received, filesize)
 
 
-def download(url: str, itag: str) -> None:
+def download(youtube: YouTube, itag: int) -> None:
     """Start downloading a YouTube video.
 
-    :param str url:
-        A valid YouTube watch URL.
+    :param YouTube youtube:
+        A valid YouTube object.
     :param str itag:
         YouTube format identifier code.
 
     """
     # 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(int(itag))
+    stream = youtube.streams.get_by_itag(itag)
     if stream is None:
-        print("Could not find a stream with itag: " + itag)
+        print("Could not find a stream with itag: {itag}".format(itag=itag))
+        print("Try one of these:")
+        display_streams(youtube)
         sys.exit()
+
+    youtube.register_on_progress_callback(on_progress)
     print("\n{fn} | {fs} bytes".format(fn=stream.default_filename, fs=stream.filesize,))
     try:
         stream.download()
@@ -171,17 +192,48 @@ def download(url: str, itag: str) -> None:
         sys.exit()
 
 
-def display_streams(url: str) -> None:
+def display_streams(youtube: YouTube) -> None:
     """Probe YouTube video and lists its available formats.
 
-    :param str url:
+    :param YouTube youtube:
         A valid YouTube watch URL.
 
     """
-    yt = YouTube(url)
-    for stream in yt.streams.all():
+    for stream in youtube.streams.all():
         print(stream)
 
 
+def _print_available_captions(captions: CaptionQuery) -> None:
+    print(
+        "Available caption codes are: {}".format(
+            ", ".join(c.code for c in captions.all())
+        )
+    )
+
+
+def download_caption(youtube: YouTube, lang_code: Optional[str]) -> None:
+    """Download a caption for the YouTube video.
+
+    :param YouTube youtube:
+        A valid YouTube object.
+    :param str lang_code:
+        Language code desired for caption file.
+        Prints available codes if the value is None
+        or the desired code is not available.
+
+    """
+    if lang_code is None:
+        _print_available_captions(youtube.captions)
+        return
+
+    caption = youtube.captions.get_by_language_code(lang_code=lang_code)
+    if caption:
+        downloaded_path = caption.download(title=youtube.title)
+        print("Saved caption file to: {}".format(downloaded_path))
+    else:
+        print("Unable to find caption with code: {}".format(lang_code))
+        _print_available_captions(youtube.captions)
+
+
 if __name__ == "__main__":
     main()

tests/test_captions.py

@@ -1,3 +1,6 @@
+from unittest import mock
+from unittest.mock import patch, mock_open
+
 from pytube import Caption, CaptionQuery
 
 
@@ -39,3 +42,46 @@ def test_caption_query_get_by_language_code_when_not_exists():
     )
     caption_query = CaptionQuery(captions=[caption1, caption2])
     assert caption_query.get_by_language_code("hello") is None
+
+
+@mock.patch("pytube.captions.Caption.generate_srt_captions")
+def test_download(srt):
+    open_mock = mock_open()
+    with patch("builtins.open", open_mock):
+        srt.return_value = ""
+        caption = Caption(
+            {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        )
+        caption.download("title")
+        assert open_mock.call_args_list[0][0][0].split("/")[-1] == "title (en).srt"
+
+
+@mock.patch("pytube.captions.Caption.generate_srt_captions")
+def test_download_with_prefix(srt):
+    open_mock = mock_open()
+    with patch("builtins.open", open_mock):
+        srt.return_value = ""
+        caption = Caption(
+            {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        )
+        caption.download("title", filename_prefix="1 ")
+        assert open_mock.call_args_list[0][0][0].split("/")[-1] == "1 title (en).srt"
+
+
+@mock.patch("pytube.captions.Caption.xml_captions")
+def test_download_xml_and_trim_extension(xml):
+    open_mock = mock_open()
+    with patch("builtins.open", open_mock):
+        xml.return_value = ""
+        caption = Caption(
+            {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        )
+        caption.download("title.xml", srt=False)
+        assert open_mock.call_args_list[0][0][0].split("/")[-1] == "title (en).xml"
+
+
+def test_repr():
+    caption = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    assert str(caption) == '<Caption lang="name1" code="en">'

tests/test_cli.py

@@ -1,13 +1,158 @@
 # -*- coding: utf-8 -*-
+import argparse
 from unittest import mock
+from unittest.mock import MagicMock, patch
 
-from pytube import cli
+import pytest
+
+from pytube import cli, StreamQuery, Caption, CaptionQuery
+
+parse_args = cli._parse_args
+
+
+@mock.patch("pytube.cli.YouTube")
+def test_download_when_itag_not_found(youtube):
+    youtube.streams = mock.Mock()
+    youtube.streams.all.return_value = []
+    youtube.streams.get_by_itag.return_value = None
+    with pytest.raises(SystemExit):
+        cli.download(youtube, 123)
+    youtube.streams.get_by_itag.assert_called_with(123)
+
+
+@mock.patch("pytube.cli.YouTube")
+@mock.patch("pytube.Stream")
+def test_download_when_itag_is_found(youtube, stream):
+    stream.itag = 123
+    youtube.streams = StreamQuery([stream])
+    with patch.object(
+        youtube.streams, "get_by_itag", wraps=youtube.streams.get_by_itag
+    ) as wrapped_itag:
+        cli.download(youtube, 123)
+        wrapped_itag.assert_called_with(123)
+    youtube.register_on_progress_callback.assert_called_with(cli.on_progress)
+    stream.download.assert_called()
 
 
 @mock.patch("pytube.cli.YouTube")
-@mock.patch("pytube.cli.sys")
-def test_download(MockYouTube, mock_sys):
-    instance = MockYouTube.return_value
-    instance.prefetch_descramble.return_value = None
-    instance.streams = mock.Mock()
-    cli.download("asdf", 123)
+@mock.patch("pytube.Stream")
+def test_display_stream(youtube, stream):
+    stream.itag = 123
+    stream.__repr__ = MagicMock(return_value="")
+    youtube.streams = StreamQuery([stream])
+    with patch.object(youtube.streams, "all", wraps=youtube.streams.all) as wrapped_all:
+        cli.display_streams(youtube)
+        wrapped_all.assert_called()
+        stream.__repr__.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube")
+def test_download_caption_with_none(youtube):
+    caption = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    youtube.captions = CaptionQuery([caption])
+    with patch.object(
+        youtube.captions, "all", wraps=youtube.captions.all
+    ) as wrapped_all:
+        cli.download_caption(youtube, None)
+        wrapped_all.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube")
+def test_download_caption_with_language_found(youtube):
+    youtube.title = "video title"
+    caption = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    caption.download = MagicMock(return_value="file_path")
+    youtube.captions = CaptionQuery([caption])
+    cli.download_caption(youtube, "en")
+    caption.download.assert_called_with(title="video title")
+
+
+@mock.patch("pytube.cli.YouTube")
+def test_download_caption_with_language_not_found(youtube):
+    caption = Caption(
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+    )
+    youtube.captions = CaptionQuery([caption])
+    with patch.object(
+        youtube.captions, "all", wraps=youtube.captions.all
+    ) as wrapped_all:
+        cli.download_caption(youtube, "blah")
+        wrapped_all.assert_called()
+
+
+@mock.patch("pytube.Stream")
+@mock.patch("io.BufferedWriter")
+def test_on_progress(stream, writer):
+    stream.filesize = 10
+    cli.display_progress_bar = MagicMock()
+    cli.on_progress(stream, "", writer, 7)
+    cli.display_progress_bar.assert_called_once_with(3, 10)
+
+
+def test_parse_args_falsey():
+    parser = argparse.ArgumentParser()
+    args = cli._parse_args(parser, ["urlhere"])
+    assert args.url == "urlhere"
+    assert args.build_playback_report is False
+    assert args.itag is None
+    assert args.list is False
+    assert args.verbosity == 0
+
+
+def test_parse_args_truthy():
+    parser = argparse.ArgumentParser()
+    args = cli._parse_args(
+        parser, ["urlhere", "--build-playback-report", "-c", "en", "-l", "--itag=10"]
+    )
+    assert args.url == "urlhere"
+    assert args.build_playback_report is True
+    assert args.itag == 10
+    assert args.list is True
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_download(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "--itag=10"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.download = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.download.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_build_playback_report(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "--build-playback-report"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.build_playback_report = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.build_playback_report.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_display_streams(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "-l"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.display_streams = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.display_streams.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_download_caption(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "-c"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.download_caption = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.download_caption.assert_called()