CaptionQuery implements Mapping from lang codes

pytube/query.py

@@ -2,7 +2,7 @@
 
 """This module provides a query interface for media streams and captions."""
 from typing import Callable, List, Optional, Union
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 
 from pytube import Stream, Caption
 from pytube.helpers import deprecated
@@ -356,7 +356,7 @@ class StreamQuery(Sequence):
         return f"{self.fmt_streams}"
 
 
-class CaptionQuery(Sequence):
+class CaptionQuery(Mapping):
     """Interface for querying the available captions."""
 
     def __init__(self, captions: List[Caption]):
@@ -366,9 +366,9 @@ class CaptionQuery(Sequence):
             list of :class:`Caption <Caption>` instances.
 
         """
-        self.captions = captions
         self.lang_code_index = {c.code: c for c in captions}
 
+    @deprecated("This object can be treated as a dictionary, i.e. captions['en']")
     def get_by_language_code(self, lang_code: str) -> Optional[Caption]:
         """Get the :class:`Caption <Caption>` for a given ``lang_code``.
 
@@ -381,20 +381,23 @@ class CaptionQuery(Sequence):
         """
         return self.lang_code_index.get(lang_code)
 
-    @deprecated("This object can be treated as a list, all() is useless")
+    @deprecated("This object can be treated as a dictionary")
     def all(self) -> List[Caption]:  # pragma: no cover
         """Get all the results represented by this query as a list.
 
         :rtype: list
 
         """
-        return self.captions
+        return list(self.lang_code_index.values())
 
-    def __getitem__(self, i: Union[slice, int]):
-        return self.captions[i]
+    def __getitem__(self, i: str):
+        return self.lang_code_index[i]
 
     def __len__(self) -> int:
-        return len(self.captions)
+        return len(self.lang_code_index)
+
+    def __iter__(self):
+        return iter(self.lang_code_index)
 
     def __repr__(self) -> str:
-        return f"{self.captions}"
+        return f"{self.lang_code_index}"

tests/test_captions.py

@@ -2,6 +2,8 @@
 from unittest import mock
 from unittest.mock import patch, mock_open, MagicMock
 
+import pytest
+
 from pytube import Caption, CaptionQuery, captions
 
 
@@ -20,9 +22,11 @@ def test_caption_query_sequence():
         {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr"}
     )
     caption_query = CaptionQuery(captions=[caption1, caption2])
-    assert caption_query.captions == [caption1, caption2]
     assert len(caption_query) == 2
-    assert caption_query[0] == caption1
+    assert caption_query["en"] == caption1
+    assert caption_query["fr"] == caption2
+    with pytest.raises(KeyError):
+        caption_query["nada"]
 
 
 def test_caption_query_get_by_language_code_when_exists():
@@ -104,7 +108,7 @@ def test_repr():
     assert str(caption) == '<Caption lang="name1" code="en">'
 
     caption_query = CaptionQuery(captions=[caption])
-    assert repr(caption_query) == '[<Caption lang="name1" code="en">]'
+    assert repr(caption_query) == '{\'en\': <Caption lang="name1" code="en">}'
 
 
 @mock.patch("pytube.request.get")