refactor StreamQuery.order_by and add tests

.idea/dictionaries/haroldmartin.xml

@@ -5,8 +5,10 @@
       <w>descrambler</w>
       <w>ficano</w>
       <w>itag</w>
+      <w>mypy</w>
       <w>noqa</w>
       <w>pytube</w>
+      <w>webm</w>
     </words>
   </dictionary>
 </component>

pytube/itags.py

@@ -95,6 +95,8 @@ HDR = [330, 331, 332, 333, 334, 335, 336, 337]
 _60FPS = [298, 299, 302, 303, 308, 315] + HDR
 _3D = [82, 83, 84, 85, 100, 101, 102]
 LIVE = [91, 92, 93, 94, 95, 96, 132, 151]
+DASH_MP4_VIDEO = [133, 134, 135, 136, 137, 138, 160, 212, 264, 266, 298, 299]
+DASH_MP4_AUDIO = [139, 140, 141, 256, 258, 325, 328]
 
 
 def get_format_profile(itag):

pytube/query.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """This module provides a query interface for media streams and captions."""
-from typing import List, Optional
+from typing import List, Optional, Dict
 
 from pytube import Stream, Caption
 
@@ -166,34 +166,45 @@ class StreamQuery:
             fmt_streams = list(filter(fn, fmt_streams))
         return StreamQuery(fmt_streams)
 
-    def order_by(self, attribute_name):
-        """Apply a sort order to a resultset.
+    def order_by(self, attribute_name: str) -> "StreamQuery":
+        """Apply a sort order to a resultset. Implicitly Filters out stream the do not have the attribute.
 
         :param str attribute_name:
             The name of the attribute to sort by.
         """
-        integer_attr_repr = {}
-        for stream in self.fmt_streams:
-            attr = getattr(stream, attribute_name)
-            if attr is None:
-                break
-            # TODO: improve this so tests can work
-            num = "".join(x for x in attr if x.isdigit())
-            integer_attr_repr[attr] = int("".join(num)) if num else None
-
-        # if every attribute has an integer representation
-        if integer_attr_repr and all(integer_attr_repr.values()):
-
-            def key(s):
-                return integer_attr_repr[getattr(s, attribute_name)]
-
+        has_attribute = [
+            s for s in self.fmt_streams if getattr(s, attribute_name) is not None
+        ]
+        integer_attr_repr: Optional[Dict[str, int]] = None
+
+        # check that the attribute value is a string
+        if len(has_attribute) > 0 and isinstance(
+            getattr(has_attribute[0], attribute_name), str
+        ):
+            # attempt to extract numerical values from string
+            try:
+                integer_attr_repr = {
+                    getattr(s, attribute_name): int(
+                        "".join(list(filter(str.isdigit, getattr(s, attribute_name))))
+                    )
+                    for s in has_attribute
+                }
+            except ValueError:
+                integer_attr_repr = None
+                pass
+
+        # lookup integer values if we have them
+        if integer_attr_repr is not None:
+            return StreamQuery(  # mypy: ignore
+                sorted(
+                    has_attribute,
+                    key=lambda s: integer_attr_repr[getattr(s, attribute_name)],  # type: ignore
+                )
+            )
         else:
-
-            def key(s):
-                return getattr(s, attribute_name)
-
-        fmt_streams = sorted(self.fmt_streams, key=key,)
-        return StreamQuery(fmt_streams)
+            return StreamQuery(
+                sorted(has_attribute, key=lambda s: getattr(s, attribute_name))
+            )
 
     def desc(self):
         """Sort streams in descending order.

tests/test_query.py

@@ -9,40 +9,38 @@ def test_count(cipher_signature):
 
 
 @pytest.mark.parametrize(
-    'test_input,expected', [
-        ({'progressive': True}, [18]),
-        ({'resolution': '720p'}, [136, 247]),
-        ({'res': '720p'}, [136, 247]),
-        ({'fps': 30, 'resolution': '480p'}, [135, 244]),
-        ({'mime_type': 'audio/mp4'}, [140]),
-        ({'type': 'audio'}, [140, 249, 250, 251]),
-        ({'subtype': '3gpp'}, []),
-        ({'abr': '128kbps'}, [140]),
-        ({'bitrate': '128kbps'}, [140]),
-        ({'audio_codec': 'opus'}, [249, 250, 251]),
-        ({'video_codec': 'vp9'}, [248, 247, 244, 243, 242, 278]),
-        ({'only_audio': True}, [140, 249, 250, 251]),
-        ({'only_video': True, 'video_codec': 'avc1.4d4015'}, [133]),
-        ({'adaptive': True, 'resolution': '1080p'}, [137, 248]),
-        ({'custom_filter_functions': [lambda s: s.itag == 18]}, [18]),
+    "test_input,expected",
+    [
+        ({"progressive": True}, [18]),
+        ({"resolution": "720p"}, [136, 247]),
+        ({"res": "720p"}, [136, 247]),
+        ({"fps": 30, "resolution": "480p"}, [135, 244]),
+        ({"mime_type": "audio/mp4"}, [140]),
+        ({"type": "audio"}, [140, 249, 250, 251]),
+        ({"subtype": "3gpp"}, []),
+        ({"abr": "128kbps"}, [140]),
+        ({"bitrate": "128kbps"}, [140]),
+        ({"audio_codec": "opus"}, [249, 250, 251]),
+        ({"video_codec": "vp9"}, [248, 247, 244, 243, 242, 278]),
+        ({"only_audio": True}, [140, 249, 250, 251]),
+        ({"only_video": True, "video_codec": "avc1.4d4015"}, [133]),
+        ({"adaptive": True, "resolution": "1080p"}, [137, 248]),
+        ({"custom_filter_functions": [lambda s: s.itag == 18]}, [18]),
     ],
 )
 def test_filters(test_input, expected, cipher_signature):
     """Ensure filters produce the expected results."""
-    result = [
-        s.itag for s
-        in cipher_signature.streams.filter(**test_input).all()
-    ]
+    result = [s.itag for s in cipher_signature.streams.filter(**test_input).all()]
     assert result == expected
 
 
-@pytest.mark.parametrize('test_input', ['first', 'last'])
+@pytest.mark.parametrize("test_input", ["first", "last"])
 def test_empty(test_input, cipher_signature):
     """Ensure :meth:`~pytube.StreamQuery.last` and
     :meth:`~pytube.StreamQuery.first` return None if the resultset is
     empty.
     """
-    query = cipher_signature.streams.filter(video_codec='vp20')
+    query = cipher_signature.streams.filter(video_codec="vp20")
     fn = getattr(query, test_input)
     assert fn() is None
 
@@ -61,78 +59,72 @@ def test_get_first(cipher_signature):
     assert cipher_signature.streams.first().itag == 18
 
 
-@pytest.mark.skip(reason="Error with int sorting")
 def test_order_by(cipher_signature):
     """Ensure :meth:`~pytube.StreamQuery.order_by` sorts the list of
     :class:`Stream <Stream>` instances in the expected order.
     """
     itags = [
-        s.itag for s in cipher_signature.streams
-        .filter(progressive=True)
-        .order_by('itag')
-        .all()
+        s.itag
+        for s in cipher_signature.streams.filter(type="audio").order_by("itag").all()
     ]
+    assert itags == [140, 249, 250, 251]
 
-    assert itags == ['17', '18', '22', '36', '43']
 
-
-@pytest.mark.skip(reason="Error with int sorting")
 def test_order_by_descending(cipher_signature):
     """Ensure :meth:`~pytube.StreamQuery.desc` sorts the list of
     :class:`Stream <Stream>` instances in the reverse order.
     """
     # numerical values
     itags = [
-        s.itag for s in cipher_signature.streams
-        .filter(progressive=True)
-        .order_by('itag')
+        s.itag
+        for s in cipher_signature.streams.filter(type="audio")
+        .order_by("itag")
         .desc()
         .all()
     ]
-    assert itags == ['43', '36', '22', '18', '17']
+    assert itags == [251, 250, 249, 140]
+
 
-    # non numerical values
+def test_order_by_non_numerical(cipher_signature):
     mime_types = [
-        s.mime_type for s in cipher_signature.streams
-        .filter(progressive=True)
-        .order_by('mime_type')
+        s.mime_type
+        for s in cipher_signature.streams.filter(res="360p")
+        .order_by("mime_type")
         .desc()
         .all()
     ]
-    assert mime_types == [
-        'video/webm', 'video/mp4',
-        'video/mp4', 'video/3gpp', 'video/3gpp',
-    ]
+    assert mime_types == ["video/webm", "video/mp4", "video/mp4"]
 
 
-@pytest.mark.skip(reason="Error with int sorting")
 def test_order_by_ascending(cipher_signature):
     """Ensure :meth:`~pytube.StreamQuery.desc` sorts the list of
     :class:`Stream <Stream>` instances in ascending order.
     """
     # numerical values
     itags = [
-        s.itag for s in cipher_signature.streams
-        .filter(progressive=True)
-        .order_by('itag')
+        s.itag
+        for s in cipher_signature.streams.filter(type="audio")
+        .order_by("itag")
         .asc()
         .all()
     ]
+    assert itags == [140, 249, 250, 251]
 
-    assert itags == ['17', '18', '22', '36', '43']
 
-    # non numerical values
+def test_order_by_non_numerical_ascending(cipher_signature):
     mime_types = [
-        s.mime_type for s in cipher_signature.streams
-        .filter(progressive=True)
-        .order_by('mime_type')
+        s.mime_type
+        for s in cipher_signature.streams.filter(res="360p")
+        .order_by("mime_type")
         .asc()
         .all()
     ]
-    assert mime_types == [
-        'video/3gpp', 'video/3gpp',
-        'video/mp4', 'video/mp4', 'video/webm',
-    ]
+    assert mime_types == ["video/mp4", "video/mp4", "video/webm"]
+
+
+def test_order_by_with_none_values(cipher_signature):
+    abrs = [s.abr for s in cipher_signature.streams.order_by("abr").asc().all()]
+    assert abrs == ["50kbps", "70kbps", "96kbps", "128kbps", "160kbps"]
 
 
 def test_get_by_itag(cipher_signature):
@@ -140,7 +132,7 @@ def test_get_by_itag(cipher_signature):
     :class:`Stream <Stream>`.
     """
     assert cipher_signature.streams.get_by_itag(18).itag == 18
-    assert cipher_signature.streams.get_by_itag('18').itag == 18
+    assert cipher_signature.streams.get_by_itag("18").itag == 18
 
 
 def test_get_by_non_existent_itag(cipher_signature):