Merge pull request #287 from Mattwmaster58/numerical

pytube/query.py

@@ -162,9 +162,23 @@ class StreamQuery:
         :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
+            num = [char for char in attr if char.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)]
+        else:
+            def key(s): return getattr(s, attribute_name)
+
         fmt_streams = sorted(
             self.fmt_streams,
-            key=lambda s: getattr(s, attribute_name),
+            key=key
         )
         return StreamQuery(fmt_streams)
 

tests/test_query.py

@@ -80,6 +80,7 @@ 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)
@@ -87,14 +88,25 @@ def test_order_by_descending(cipher_signature):
         .desc()
         .all()
     ]
-
     assert itags == ['43', '36', '22', '18', '17']
 
+    # non numerical values
+    mime_types = [
+        s.mime_type for s in cipher_signature.streams
+        .filter(progressive=True)
+        .order_by('mime_type')
+        .desc()
+        .all()
+    ]
+    assert mime_types == ['video/webm', 'video/mp4',
+                          'video/mp4', 'video/3gpp', 'video/3gpp']
+
 
 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)
@@ -105,6 +117,17 @@ def test_order_by_ascending(cipher_signature):
 
     assert itags == ['17', '18', '22', '36', '43']
 
+    # non numerical values
+    mime_types = [
+        s.mime_type for s in cipher_signature.streams
+        .filter(progressive=True)
+        .order_by('mime_type')
+        .asc()
+        .all()
+    ]
+    assert mime_types == ['video/3gpp', 'video/3gpp',
+                          'video/mp4', 'video/mp4', 'video/webm']
+
 
 def test_get_by_itag(cipher_signature):
     """Ensure :meth:`~pytube.StreamQuery.get_by_itag` returns the expected