Improved caption selection, and updated tests to reflect this. (#783)

pytube/captions.py

@@ -23,7 +23,12 @@ class Caption:
         """
         self.url = caption_track.get("baseUrl")
         self.name = caption_track["name"]["simpleText"]
-        self.code = caption_track["languageCode"]
+        # Use "vssId" instead of "languageCode", fix issue #779
+        self.code = caption_track["vssId"]
+        # Remove preceding '.' for backwards compatibility, e.g.:
+        # English -> vssId: .en, languageCode: en
+        # English (auto-generated) -> vssId: a.en, languageCode: en
+        self.code = self.code.strip('.')
 
     @property
     def xml_captions(self) -> str:

tests/test_captions.py

@@ -14,17 +14,17 @@ from pytube import captions
 
 def test_float_to_srt_time_format():
     caption1 = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     assert caption1.float_to_srt_time_format(3.89) == "00:00:03,890"
 
 
 def test_caption_query_sequence():
     caption1 = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     caption2 = Caption(
-        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
+        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr", "vssId": ".fr"}
     )
     caption_query = CaptionQuery(captions=[caption1, caption2])
     assert len(caption_query) == 2
@@ -36,10 +36,10 @@ def test_caption_query_sequence():
 
 def test_caption_query_get_by_language_code_when_exists():
     caption1 = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     caption2 = Caption(
-        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
+        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr", "vssId": ".fr"}
     )
     caption_query = CaptionQuery(captions=[caption1, caption2])
     assert caption_query["en"] == caption1
@@ -47,10 +47,10 @@ def test_caption_query_get_by_language_code_when_exists():
 
 def test_caption_query_get_by_language_code_when_not_exists():
     caption1 = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     caption2 = Caption(
-        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
+        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr", "vssId": ".fr"}
     )
     caption_query = CaptionQuery(captions=[caption1, caption2])
     with pytest.raises(KeyError):
@@ -68,6 +68,7 @@ def test_download(srt):
                 "url": "url1",
                 "name": {"simpleText": "name1"},
                 "languageCode": "en",
+                "vssId": ".en"
             }
         )
         caption.download("title")
@@ -86,6 +87,7 @@ def test_download_with_prefix(srt):
                 "url": "url1",
                 "name": {"simpleText": "name1"},
                 "languageCode": "en",
+                "vssId": ".en"
             }
         )
         caption.download("title", filename_prefix="1 ")
@@ -106,6 +108,7 @@ def test_download_with_output_path(srt):
                 "url": "url1",
                 "name": {"simpleText": "name1"},
                 "languageCode": "en",
+                "vssId": ".en"
             }
         )
         file_path = caption.download("title", output_path="blah")
@@ -123,6 +126,7 @@ def test_download_xml_and_trim_extension(xml):
                 "url": "url1",
                 "name": {"simpleText": "name1"},
                 "languageCode": "en",
+                "vssId": ".en"
             }
         )
         caption.download("title.xml", srt=False)
@@ -133,7 +137,7 @@ def test_download_xml_and_trim_extension(xml):
 
 def test_repr():
     caption = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     assert str(caption) == '<Caption lang="name1" code="en">'
 
@@ -145,7 +149,7 @@ def test_repr():
 def test_xml_captions(request_get):
     request_get.return_value = "test"
     caption = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     assert caption.xml_captions == "test"
 
@@ -158,7 +162,7 @@ def test_generate_srt_captions(request):
         "如要啓動字幕，請按一下這裡的圖示。</text></transcript>"
     )
     caption = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     assert caption.generate_srt_captions() == (
         "1\n"

tests/test_cli.py

@@ -71,7 +71,7 @@ def test_display_stream(youtube, stream):
 def test_download_caption_with_none(youtube, print_available):
     # Given
     caption = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     youtube.captions = CaptionQuery([caption])
     # When
@@ -84,7 +84,7 @@ def test_download_caption_with_none(youtube, print_available):
 def test_download_caption_with_language_found(youtube):
     youtube.title = "video title"
     caption = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     caption.download = MagicMock(return_value="file_path")
     youtube.captions = CaptionQuery([caption])
@@ -97,7 +97,7 @@ def test_download_caption_with_language_found(youtube):
 def test_download_caption_with_lang_not_found(youtube, print_available):
     # Given
     caption = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     youtube.captions = CaptionQuery([caption])
     # When
@@ -109,10 +109,10 @@ def test_download_caption_with_lang_not_found(youtube, print_available):
 def test_print_available_captions(capsys):
     # Given
     caption1 = Caption(
-        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
+        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
     )
     caption2 = Caption(
-        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
+        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr", "vssId": ".fr"}
     )
     query = CaptionQuery([caption1, caption2])
     # When