diff --git a/scripts/wk/std.py b/scripts/wk/std.py index 0d0989a9..3f866545 100644 --- a/scripts/wk/std.py +++ b/scripts/wk/std.py @@ -27,6 +27,9 @@ COLORS = { 'CYAN': '\033[36m', } LOG = logging.getLogger(__name__) +REGEX_SIZE_STRING = re.compile( + r'(?P\d+\.?\d*)\s+(?P[PTGMKB])(?PI?)B?' + ) # Functions @@ -187,6 +190,42 @@ def sleep(seconds=2): time.sleep(seconds) +def string_to_bytes(size, assume_binary=False): + """Convert human-readable size str to bytes and return an int.""" + LOG.debug('size: %s, assume_binary: %s', size, assume_binary) + scale = 1000 + size = str(size) + tmp = REGEX_SIZE_STRING.search(size.upper()) + + # Set scale + if tmp.group('binary') or assume_binary: + scale = 1024 + + # Convert to bytes + if tmp: + size = float(tmp.group('size')) + units = tmp.group('units') + if units == 'P': + size *= scale ** 5 + if units == 'T': + size *= scale ** 4 + elif units == 'G': + size *= scale ** 3 + elif units == 'M': + size *= scale ** 2 + elif units == 'K': + size *= scale ** 1 + elif units == 'B': + size *= scale ** 0 + size = int(size) + else: + return -1 + + # Done + LOG.debug('bytes: %s', size) + return size + + def strip_colors(string): """Strip known ANSI color escapes from string, returns str.""" LOG.debug('string: %s', string)