diff --git a/gallery_dl/util.py b/gallery_dl/util.py index 03c7751b..3cf3000c 100644 --- a/gallery_dl/util.py +++ b/gallery_dl/util.py @@ -84,6 +84,22 @@ def bdecode(data, alphabet="0123456789"): return num +def parse_bytes(value, suffixes="bkmgtp"): + """Convert a bytes-amount ("500k", "2.5M", ...) to int""" + last = value[-1].lower() + + if last in suffixes: + mul = 1024 ** suffixes.index(last) + value = value[:-1] + else: + mul = 1 + + try: + return round(float(value) * mul) + except ValueError: + return 0 + + def combine_dict(a, b): """Recursively combine the contents of b into a""" for key, value in b.items(): diff --git a/test/test_util.py b/test/test_util.py index 690bd288..137a0f8d 100644 --- a/test/test_util.py +++ b/test/test_util.py @@ -189,6 +189,22 @@ class TestOther(unittest.TestCase): self.assertEqual(util.bdecode("1111011", "01"), 123) self.assertEqual(util.bdecode("AAAABAA", "BA"), 123) + def test_parse_bytes(self): + self.assertEqual(util.parse_bytes("50"), 50) + self.assertEqual(util.parse_bytes("50k"), 50 * 1024**1) + self.assertEqual(util.parse_bytes("50m"), 50 * 1024**2) + self.assertEqual(util.parse_bytes("50g"), 50 * 1024**3) + self.assertEqual(util.parse_bytes("50t"), 50 * 1024**4) + self.assertEqual(util.parse_bytes("50p"), 50 * 1024**5) + + self.assertEqual(util.parse_bytes("123.456"), 123) + self.assertEqual(util.parse_bytes("123.567"), 124) + self.assertEqual(util.parse_bytes("0.5M"), round(0.5 * 1024**2)) + + self.assertEqual(util.parse_bytes("NaN"), 0) + self.assertEqual(util.parse_bytes("invalid"), 0) + self.assertEqual(util.parse_bytes(" 123 kb "), 0) + def test_combine_dict(self): self.assertEqual( util.combine_dict({}, {}),