[utils] Improve parse_count

dirkf · Dec 24, 2021 · 352d5da · 352d5da
1 parent d43de68
commit 352d5da
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -1156,9 +1156,16 @@ def test_parse_count(self):
         self.assertEqual(parse_count('1000'), 1000)
         self.assertEqual(parse_count('1.000'), 1000)
         self.assertEqual(parse_count('1.1k'), 1100)
+        self.assertEqual(parse_count('1.1 k'), 1100)
+        self.assertEqual(parse_count('1,1 k'), 1100)
         self.assertEqual(parse_count('1.1kk'), 1100000)
         self.assertEqual(parse_count('1.1kk '), 1100000)
+        self.assertEqual(parse_count('1,1kk'), 1100000)
+        self.assertEqual(parse_count('100 views'), 100)
+        self.assertEqual(parse_count('1,100 views'), 1100)
         self.assertEqual(parse_count('1.1kk views'), 1100000)
+        self.assertEqual(parse_count('10M views'), 10000000)
+        self.assertEqual(parse_count('has 10M views'), 10000000)
 
     def test_parse_resolution(self):
         self.assertEqual(parse_resolution(None), {})

diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py
@@ -2211,7 +2211,7 @@ def parse_count(s):
     if s is None:
         return None
 
-    s = s.strip()
+    s = re.sub(r'^[^\d]+\s', '', s).strip()
 
     if re.match(r'^[\d,.]+$', s):
         return str_to_int(s)
@@ -2223,9 +2223,17 @@ def parse_count(s):
         'M': 1000 ** 2,
         'kk': 1000 ** 2,
         'KK': 1000 ** 2,
+        'b': 1000 ** 3,
+        'B': 1000 ** 3,
     }
 
-    return lookup_unit_table(_UNIT_TABLE, s)
+    ret = lookup_unit_table(_UNIT_TABLE, s)
+    if ret is not None:
+        return ret
+
+    mobj = re.match(r'([\d,.]+)(?:$|\s)', s)
+    if mobj:
+        return str_to_int(mobj.group(1))
 
 
 def parse_resolution(s):