Merge pull request #106 from HansBug/dev/size

dev(narugo): add new size params

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)

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')