From 64c464a144e2a96ec21a717d191217edda9107a4 Mon Sep 17 00:00:00 2001 From: pukkandan Date: Thu, 17 Nov 2022 08:40:34 +0530 Subject: [PATCH] [utils] Move `FileDownloader.parse_bytes` into utils --- yt_dlp/__init__.py | 18 +++++++++--------- yt_dlp/downloader/common.py | 9 ++------- yt_dlp/utils.py | 19 ++++++++++++++----- 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/yt_dlp/__init__.py b/yt_dlp/__init__.py index 726fb0685c..c03e6e6915 100644 --- a/yt_dlp/__init__.py +++ b/yt_dlp/__init__.py @@ -16,7 +16,6 @@ import sys from .compat import compat_shlex_quote from .cookies import SUPPORTED_BROWSERS, SUPPORTED_KEYRINGS -from .downloader import FileDownloader from .downloader.external import get_external_downloader from .extractor import list_extractor_classes from .extractor.adobepass import MSO_INFO @@ -50,6 +49,7 @@ from .utils import ( format_field, int_or_none, match_filter_func, + parse_bytes, parse_duration, preferredencoding, read_batch_urls, @@ -281,19 +281,19 @@ def validate_options(opts): raise ValueError(f'invalid {key} retry sleep expression {expr!r}') # Bytes - def parse_bytes(name, value): + def validate_bytes(name, value): if value is None: return None - numeric_limit = FileDownloader.parse_bytes(value) + numeric_limit = parse_bytes(value) validate(numeric_limit is not None, 'rate limit', value) return numeric_limit - opts.ratelimit = parse_bytes('rate limit', opts.ratelimit) - opts.throttledratelimit = parse_bytes('throttled rate limit', opts.throttledratelimit) - opts.min_filesize = parse_bytes('min filesize', opts.min_filesize) - opts.max_filesize = parse_bytes('max filesize', opts.max_filesize) - opts.buffersize = parse_bytes('buffer size', opts.buffersize) - opts.http_chunk_size = parse_bytes('http chunk size', opts.http_chunk_size) + opts.ratelimit = validate_bytes('rate limit', opts.ratelimit) + opts.throttledratelimit = validate_bytes('throttled rate limit', opts.throttledratelimit) + opts.min_filesize = validate_bytes('min filesize', opts.min_filesize) + opts.max_filesize = validate_bytes('max filesize', opts.max_filesize) + opts.buffersize = validate_bytes('buffer size', opts.buffersize) + opts.http_chunk_size = validate_bytes('http chunk size', opts.http_chunk_size) # Output templates def validate_outtmpl(tmpl, msg): diff --git a/yt_dlp/downloader/common.py b/yt_dlp/downloader/common.py index 8d110c3747..fe36332506 100644 --- a/yt_dlp/downloader/common.py +++ b/yt_dlp/downloader/common.py @@ -15,7 +15,6 @@ from ..minicurses import ( from ..utils import ( IDENTITY, NO_DEFAULT, - NUMBER_RE, LockingUnsupportedError, Namespace, RetryManager, @@ -24,6 +23,7 @@ from ..utils import ( encodeFilename, format_bytes, join_nonempty, + parse_bytes, remove_start, sanitize_open, shell_quote, @@ -180,12 +180,7 @@ class FileDownloader: @staticmethod def parse_bytes(bytestr): """Parse a string indicating a byte quantity into an integer.""" - matchobj = re.match(rf'(?i)^({NUMBER_RE})([kMGTPEZY]?)$', bytestr) - if matchobj is None: - return None - number = float(matchobj.group(1)) - multiplier = 1024.0 ** 'bkmgtpezy'.index(matchobj.group(2).lower()) - return int(round(number * multiplier)) + parse_bytes(bytestr) def slow_down(self, start_time, now, byte_counter): """Sleep if the download speed is over the rate limit.""" diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py index 7cba13678f..9b6977b6df 100644 --- a/yt_dlp/utils.py +++ b/yt_dlp/utils.py @@ -2289,15 +2289,24 @@ def format_bytes(bytes): return format_decimal_suffix(bytes, '%.2f%sB', factor=1024) or 'N/A' -def lookup_unit_table(unit_table, s): +def lookup_unit_table(unit_table, s, strict=False): + num_re = NUMBER_RE if strict else NUMBER_RE.replace(R'\.', '[,.]') units_re = '|'.join(re.escape(u) for u in unit_table) - m = re.match( - r'(?P[0-9]+(?:[,.][0-9]*)?)\s*(?P%s)\b' % units_re, s) + m = (re.fullmatch if strict else re.match)( + rf'(?P{num_re})\s*(?P{units_re})\b', s) if not m: return None - num_str = m.group('num').replace(',', '.') + + num = float(m.group('num').replace(',', '.')) mult = unit_table[m.group('unit')] - return int(float(num_str) * mult) + return round(num * mult) + + +def parse_bytes(s): + """Parse a string indicating a byte quantity into an integer""" + return lookup_unit_table( + {u: 1024**i for i, u in enumerate(['', *'KMGTPEZY'])}, + s.upper(), strict=True) def parse_filesize(s):