diff --git a/hbutils/scale/size.py b/hbutils/scale/size.py index 63e6f8ea9b7..d6d1d456db3 100644 --- a/hbutils/scale/size.py +++ b/hbutils/scale/size.py @@ -1,10 +1,12 @@ """ Overview: - Useful utilities for memory size units, such as MB/KB/B. + This module provides useful utilities for handling memory size units, such as MB/KB/B. + It includes functions for converting various size representations to bytes and formatting + size values as human-readable strings. """ import warnings from enum import IntEnum, unique -from typing import Union, Optional +from typing import Union, Optional, Literal from bitmath import Byte, NIST, SI from bitmath import parse_string_unsafe as parse_bytes @@ -19,6 +21,23 @@ def _is_int(value: Union[int, float], stacklevel: int = 4) -> int: + """ + Convert a value to an integer, with a warning if rounding occurs. + + :param value: The value to convert to an integer. + :type value: Union[int, float] + :param stacklevel: The stack level for the warning, default is 4. + :type stacklevel: int + + :return: The integer representation of the value. + :rtype: int + + :raises AssertionError: If the input is neither float nor int. + + This function is used internally to ensure that byte values are always integers. + If a float is provided, it will be rounded to the nearest integer, and a warning + will be issued if the rounding causes a significant change in the value. + """ if isinstance(value, float): _actual = int(round(value)) _delta = abs(value - _actual) @@ -34,6 +53,21 @@ def _is_int(value: Union[int, float], stacklevel: int = 4) -> int: def _base_size_to_bytes(size, stacklevel: int = 4) -> int: + """ + Convert various size representations to bytes. + + :param size: The size to convert, can be int, float, str, or Byte object. + :type size: Union[int, float, str, Byte] + :param stacklevel: The stack level for warnings, default is 4. + :type stacklevel: int + + :return: The size in bytes as an integer. + :rtype: int + + :raises TypeError: If the input type is not supported. + + This function is used internally to handle different input types and convert them to bytes. + """ if isinstance(size, (float, int)): return _is_int(size, stacklevel) elif isinstance(size, str): @@ -54,16 +88,20 @@ def _base_size_to_bytes(size, stacklevel: int = 4) -> int: def size_to_bytes(size: _SIZE_TYPING) -> int: """ - Overview: - Turn any types of memory size to integer value in bytes. + Convert any type of memory size representation to an integer value in bytes. + + :param size: Any type of size information. + :type size: Union[int, float, str, Byte] - Arguments: - - size (:obj:`Union[int, float, str, Byte]`): Any types of size information. + :return: Size in bytes as an integer. + :rtype: int - Returns: - - bytes (:obj:`int`): Int formatted size in bytes. + This function can handle various input types: + * Integers and floats are treated as byte values. + * Strings are parsed as size representations (e.g., "23356 KB"). + * Byte objects from the bitmath library are converted to their byte value. - Examples:: + Examples: >>> from hbutils.scale import size_to_bytes >>> size_to_bytes(23344) 23344 @@ -81,22 +119,40 @@ def size_to_bytes(size: _SIZE_TYPING) -> int: @int_enum_loads(name_preprocess=str.upper) @unique class SizeSystem(IntEnum): + """ + Enumeration of size systems used for formatting. + + NIST: Uses binary prefixes (KiB, MiB, GiB, etc.) + SI: Uses decimal prefixes (KB, MB, GB, etc.) + """ NIST = NIST SI = SI -def size_to_bytes_str(size: _SIZE_TYPING, precision: Optional[int] = None, system='nist') -> str: +def size_to_bytes_str(size: _SIZE_TYPING, precision: Optional[int] = None, + sigfigs: Optional[int] = None, + system: Literal['nist', 'si'] = 'nist') -> str: """ - Overview: - Turn any types of memory size to string value in the best unit. + Convert any type of memory size to a string value in the most appropriate unit. + + :param size: Any type of size information. + :type size: Union[int, float, str, Byte] + :param precision: Precision for float values. Default is None, which shows the original float number. + :type precision: Optional[int] + :param sigfigs: Number of significant figures to use. If specified, overrides precision. + :type sigfigs: Optional[int] + :param system: The unit system to use, either ``nist`` (binary) or ``si`` (decimal). Default is ``nist``. + :type system: Literal['nist', 'si'] - :param size: Any types of size information. - :param precision: Precsion for float values. Default is ``None`` which means just show the original float number. + :return: String formatted size value in the best unit. + :rtype: str - Returns: - - bytes (:obj:`int`): String formatted size value in the best unit. + This function provides a human-readable representation of size values. It automatically + chooses the most appropriate unit (B, KB/KiB, MB/MiB, etc.) based on the size. - Examples:: + The ``system`` parameter allows choosing between NIST (binary, e.g., KiB) and SI (decimal, e.g., KB) units. + + Examples: >>> from hbutils.scale import size_to_bytes_str >>> size_to_bytes_str(23344) '22.796875 KiB' @@ -116,9 +172,14 @@ def size_to_bytes_str(size: _SIZE_TYPING, precision: Optional[int] = None, syste >>> size_to_bytes_str('3.54 GB', precision=3, system='si') '3.540 GB' """ - system = SizeSystem.loads(system) - if precision is None: - format_str = "{value} {unit}" - else: + try: + system = SizeSystem.loads(system) + except KeyError: + raise ValueError(f'Invalid System Type - {system!r}.') + if sigfigs is not None: + format_str = f"{{value:.{sigfigs}g}} {{unit}}" + elif precision is not None: format_str = f"{{value:.{precision}f}} {{unit}}" + else: + format_str = "{value} {unit}" return Byte(_base_size_to_bytes(size)).best_prefix(system.value).format(format_str) diff --git a/test/scale/test_size.py b/test/scale/test_size.py index 8e901d62f68..420b0e04c38 100644 --- a/test/scale/test_size.py +++ b/test/scale/test_size.py @@ -31,3 +31,36 @@ def test_size_to_bytes_str(self): assert size_to_bytes_str('3.54 GB', precision=3) == '3.297 GiB' assert size_to_bytes_str('3.54 GB', system='si') == '3.54 GB' assert size_to_bytes_str('3.54 GB', system='si', precision=3) == '3.540 GB' + + def test_size_to_bytes_with_int(self): + assert size_to_bytes(1024) == 1024 + + def test_size_to_bytes_with_float(self): + assert size_to_bytes(1024.0) == 1024 + + def test_size_to_bytes_with_str(self): + assert size_to_bytes('1 KB') == 1000 + + def test_size_to_bytes_with_unsupported_type(self): + with pytest.raises(TypeError): + size_to_bytes([1024]) + + def test_size_to_bytes_str(self): + assert size_to_bytes_str(1024) == '1.0 KiB' + + def test_size_to_bytes_str_with_precision(self): + assert size_to_bytes_str(1500, precision=2) == '1.46 KiB' + + def test_size_to_bytes_str_with_sigfigs(self): + assert size_to_bytes_str(1500, sigfigs=3) == '1.46 KiB' + + def test_size_to_bytes_str_with_system(self): + assert size_to_bytes_str(1000, system='si') == '1.0 kB' + + def test_size_to_bytes_str_with_invalid_system(self): + with pytest.raises(ValueError): + size_to_bytes_str(1000, system='invalid') + + def test_size_to_bytes_str_warning(self): + with pytest.warns(UserWarning): + size_to_bytes_str('3.54 GiB')