diff options
| -rw-r--r-- | test/test_utils.py | 10 | ||||
| -rw-r--r-- | youtube_dl/utils.py | 37 |
2 files changed, 40 insertions, 7 deletions
diff --git a/test/test_utils.py b/test/test_utils.py index 2bcf8ecf0..bc28ceb34 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -41,6 +41,7 @@ from youtube_dl.utils import ( orderedSet, parse_duration, parse_filesize, + parse_count, parse_iso8601, read_batch_urls, sanitize_filename, @@ -653,6 +654,15 @@ class TestUtil(unittest.TestCase): self.assertEqual(parse_filesize('1.2Tb'), 1200000000000) self.assertEqual(parse_filesize('1,24 KB'), 1240) + def test_parse_count(self): + self.assertEqual(parse_count(None), None) + self.assertEqual(parse_count(''), None) + self.assertEqual(parse_count('0'), 0) + 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.1kk'), 1100000) + def test_version_tuple(self): self.assertEqual(version_tuple('1'), (1,)) self.assertEqual(version_tuple('10.23.344'), (10, 23, 344)) diff --git a/youtube_dl/utils.py b/youtube_dl/utils.py index 22a39a0ab..1d6ee9964 100644 --- a/youtube_dl/utils.py +++ b/youtube_dl/utils.py @@ -1316,6 +1316,17 @@ def format_bytes(bytes): return '%.2f%s' % (converted, suffix) +def lookup_unit_table(unit_table, s): + units_re = '|'.join(re.escape(u) for u in unit_table) + m = re.match( + r'(?P<num>[0-9]+(?:[,.][0-9]*)?)\s*(?P<unit>%s)' % units_re, s) + if not m: + return None + num_str = m.group('num').replace(',', '.') + mult = unit_table[m.group('unit')] + return int(float(num_str) * mult) + + def parse_filesize(s): if s is None: return None @@ -1359,16 +1370,28 @@ def parse_filesize(s): 'Yb': 1000 ** 8, } - units_re = '|'.join(re.escape(u) for u in _UNIT_TABLE) - m = re.match( - r'(?P<num>[0-9]+(?:[,.][0-9]*)?)\s*(?P<unit>%s)' % units_re, s) - if not m: + return lookup_unit_table(_UNIT_TABLE, s) + + +def parse_count(s): + if s is None: return None - num_str = m.group('num').replace(',', '.') - mult = _UNIT_TABLE[m.group('unit')] - return int(float(num_str) * mult) + s = s.strip() + + if re.match(r'^[\d,.]+$', s): + return str_to_int(s) + + _UNIT_TABLE = { + 'k': 1000, + 'K': 1000, + 'm': 1000 ** 2, + 'M': 1000 ** 2, + 'kk': 1000 ** 2, + 'KK': 1000 ** 2, + } + return lookup_unit_table(_UNIT_TABLE, s) def month_by_name(name): """ Return the number of a month by (locale-independently) English name """ |