⚡️ Speed up function size_str by 24%

[codeflash-ai]

📄 24% (0.24x) speedup for size_str in src/datasets/utils/py_utils.py

⏱️ Runtime : 1.61 milliseconds → 1.30 milliseconds (best of 71 runs)

📝 Explanation and details

The optimized code achieves a 23% runtime improvement by eliminating repeated allocations and type conversions that occurred on every function call.

Key Optimizations

1. Module-level constant declaration
The unit conversion table _NAME_LIST is moved from inside the function to module scope. This prevents Python from rebuilding a 5-element list on every invocation - a pure overhead since the conversion factors never change.

2. Pre-computed float conversions
The size values (250, 240, etc.) are converted to floats at module initialization rather than during the division operation. This eliminates integer-to-float conversions inside the hot loop where divisions occur.

3. Tuple instead of list
Using a tuple rather than a list provides a minor benefit since tuples are immutable and slightly more efficient to iterate over.

Performance Impact

Based on the annotated tests, the optimization shows consistent gains across all test scenarios:

• Small byte values (< 1 KiB): 25-27% faster
• Unit boundary tests (exact KiB/MiB/GiB): 8-19% faster
• Large-scale tests with 100-1000 iterations: 20-29% faster
• Edge cases (negative values, fractional sizes): 16-43% faster

The speedup is most pronounced when the function is called repeatedly, as seen in test_large_scale_many_byte_values (28.7% faster over 1000 iterations).

Real-World Context

From function_references, this function is called during dataset download/preparation in builder.py, specifically when checking disk space and logging download progress. The logging calls occur:

1. When checking available disk space before generation
2. When reporting download/generation sizes to the user
3. After successful dataset preparation

While not in a tight loop, these operations occur during dataset initialization - a user-facing workflow where perceived performance matters. The 23% reduction means faster feedback during dataset downloads, particularly for workflows that process multiple datasets sequentially.

The optimization is universally beneficial across input sizes and types with no trade-offs in correctness or behavior.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 1649 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import pytest  # used for our unit tests
from src.datasets.utils.py_utils import size_str

def _expected_size_str_reference(size_in_bytes):
    """
    Independent helper implementing the same intended algorithm to compute the expected string.
    This is intentionally simple and mirrors the documented behavior (unit list and thresholds).
    If the implementation under test is mutated (e.g., different thresholds or unit order),
    differences should cause test failures.
    """
    if not size_in_bytes:
        return "Unknown size"
    name_list = [("PiB", 2**50), ("TiB", 2**40), ("GiB", 2**30), ("MiB", 2**20), ("KiB", 2**10)]
    size_as_float = float(size_in_bytes)
    for name, bytes_count in name_list:
        val = size_as_float / bytes_count
        if val >= 1.0:
            return f"{val:.2f} {name}"
    return f"{int(size_as_float)} bytes"

def test_none_and_zero_return_unknown():
    # None should be considered unknown
    codeflash_output = size_str(None) # 704ns -> 693ns (1.59% faster)
    # 0 is falsy and should be considered unknown
    codeflash_output = size_str(0) # 250ns -> 263ns (4.94% slower)
    # Empty string is falsy; float('') would raise, but the check happens first so it should be "Unknown size"
    # Note: passing '' to the function returns "Unknown size" because if not size_in_bytes is True.
    codeflash_output = size_str("") # 214ns -> 215ns (0.465% slower)
    # False is falsy -> Unknown size
    codeflash_output = size_str(False) # 200ns -> 196ns (2.04% faster)

@pytest.mark.parametrize(
    "input_bytes, expected",
    [
        # Basic KiB example: 1536 bytes = 1.5 KiB -> formatted with two decimals
        (1536, "1.50 KiB"),
        # Basic MiB example: 1.5 * 2**20 bytes
        (int(1.5 * (2 ** 20)), "1.50 MiB"),
        # Basic GiB example: 1.5 * 2**30 bytes
        (1.5 * (2 ** 30), "1.50 GiB"),
        # Exact KiB boundary
        (2 ** 10, "1.00 KiB"),
        # Exact MiB boundary
        (2 ** 20, "1.00 MiB"),
        # Exact GiB boundary
        (2 ** 30, "1.00 GiB"),
        # Exact TiB boundary
        (2 ** 40, "1.00 TiB"),
        # PiB example (larger unit)
        (3 * (2 ** 50), "3.00 PiB"),
    ],
)
def test_basic_unit_conversions(input_bytes, expected):
    # For a variety of common sizes ensure the expected human-readable string is returned.
    codeflash_output = size_str(input_bytes) # 34.1μs -> 29.5μs (15.7% faster)

def test_boundary_values_around_kib_and_gib():
    # Just below 1 KiB should be expressed in bytes
    codeflash_output = size_str(1023) # 3.13μs -> 2.58μs (21.5% faster)
    # Exactly 1 KiB should be reported as 1.00 KiB
    codeflash_output = size_str(1024) # 2.73μs -> 2.48μs (10.4% faster)
    # Just above 1 KiB should be formatted in KiB and still round to two decimals
    codeflash_output = size_str(1025) # 1.26μs -> 1.01μs (24.4% faster)
    # Just below 1 GiB should be expressed in bytes
    codeflash_output = size_str((2 ** 30) - 1) # 1.16μs -> 947ns (22.2% faster)
    # Exactly 1 GiB
    codeflash_output = size_str(2 ** 30) # 1.01μs -> 1.03μs (1.55% slower)

def test_non_integer_and_boolean_inputs():
    # String numeric inputs that can be converted to float are accepted
    codeflash_output = size_str("2048") # 4.92μs -> 4.20μs (17.0% faster)
    # True is truthy -> float(True) == 1.0 -> no unit matches -> "1 bytes"
    codeflash_output = size_str(True) # 1.67μs -> 1.44μs (16.3% faster)
    # Ensure invalid numeric strings raise the underlying float conversion error
    with pytest.raises(ValueError):
        size_str("not-a-number") # 2.60μs -> 2.41μs (7.80% faster)

def test_negative_values_and_signed_behavior():
    # Negative numbers are truthy (nonzero) and will not match any >=1.0 thresholds,
    # so the function should ultimately return the integer representation with "bytes"
    codeflash_output = size_str(-1024) # 3.24μs -> 2.62μs (23.8% faster)

def test_small_fractional_values_less_than_one_byte():
    # Values strictly between 0 and 1 are truthy when not equal to 0, but they become float and
    # end up as "0 bytes" after int conversion in the fallback.
    codeflash_output = size_str(0.9) # 2.96μs -> 2.33μs (27.2% faster)
    codeflash_output = size_str(0.9999) # 1.18μs -> 822ns (42.9% faster)

def test_rounding_behavior_near_unit_thresholds():
    # Slightly less than 2 GiB should still format with two decimals and may round up to 2.00 GiB
    almost_2_gib = (2 ** 30) * 1.9999
    # The formatting to two decimals may round 1.9999 -> 2.00
    codeflash_output = size_str(almost_2_gib); result = codeflash_output # 4.07μs -> 3.38μs (20.6% faster)

def test_large_scale_varied_inputs():
    # Large-scale test with a diverse but bounded number of samples (100 samples).
    # This assesses consistent behavior across a range from bytes up to PiB.
    samples = []
    # create 100 diverse sample sizes across ranges without using loops >1000 iterations
    for i in range(100):
        # Spread values logarithmically across many orders of magnitude:
        # Use powers from 0..50 scaled by fraction to generate variety
        power = (i % 51)  # cycles through 0..50 (51 distinct exponents)
        multiplier = 1 + (i / 100.0)  # slightly vary the multiplier to create non-exact boundaries
        val = int((2 ** power) * multiplier)
        # Avoid creating extremely large numbers beyond PiB*10 to remain practical
        if val == 0:
            val = i + 1
        samples.append(val)

    # Check each sample against the reference implementation above to detect even small divergences.
    for s in samples:
        codeflash_output = size_str(s) # 93.9μs -> 76.8μs (22.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from src.datasets.utils.py_utils import size_str

def test_basic_bytes_under_kib():
    """Test formatting of sizes less than 1 KiB (should return raw bytes)."""
    codeflash_output = size_str(512) # 3.23μs -> 2.55μs (26.7% faster)

def test_basic_one_byte():
    """Test formatting of exactly 1 byte."""
    codeflash_output = size_str(1) # 3.05μs -> 2.42μs (25.9% faster)

def test_basic_kib_boundary():
    """Test formatting at exactly 1 KiB."""
    codeflash_output = size_str(2**10) # 4.43μs -> 3.72μs (19.2% faster)

def test_basic_mib_boundary():
    """Test formatting at exactly 1 MiB."""
    codeflash_output = size_str(2**20) # 4.36μs -> 3.69μs (18.1% faster)

def test_basic_gib_boundary():
    """Test formatting at exactly 1 GiB."""
    codeflash_output = size_str(2**30) # 4.13μs -> 3.82μs (8.14% faster)

def test_basic_tib_boundary():
    """Test formatting at exactly 1 TiB."""
    codeflash_output = size_str(2**40) # 4.02μs -> 3.68μs (9.32% faster)

def test_basic_pib_boundary():
    """Test formatting at exactly 1 PiB."""
    codeflash_output = size_str(2**50) # 3.91μs -> 3.43μs (14.1% faster)

def test_basic_kib_with_fraction():
    """Test formatting of KiB with fractional part."""
    codeflash_output = size_str(int(1.5 * 2**10)); result = codeflash_output # 4.34μs -> 3.71μs (17.2% faster)

def test_basic_mib_with_fraction():
    """Test formatting of MiB with fractional part."""
    codeflash_output = size_str(int(1.5 * 2**20)); result = codeflash_output # 4.39μs -> 3.72μs (18.0% faster)

def test_basic_gib_with_fraction():
    """Test formatting of GiB with fractional part."""
    codeflash_output = size_str(int(1.5 * 2**30)); result = codeflash_output # 4.30μs -> 3.82μs (12.5% faster)

def test_basic_tib_with_fraction():
    """Test formatting of TiB with fractional part."""
    codeflash_output = size_str(int(1.5 * 2**40)); result = codeflash_output # 4.05μs -> 3.75μs (8.04% faster)

def test_basic_pib_with_fraction():
    """Test formatting of PiB with fractional part."""
    codeflash_output = size_str(int(1.5 * 2**50)); result = codeflash_output # 3.90μs -> 3.43μs (13.8% faster)

def test_basic_float_input_kib():
    """Test that float inputs are handled correctly for KiB."""
    codeflash_output = size_str(2048.0) # 4.34μs -> 3.51μs (23.6% faster)

def test_basic_float_input_mib():
    """Test that float inputs are handled correctly for MiB."""
    codeflash_output = size_str(2.5 * 2**20); result = codeflash_output # 4.24μs -> 3.65μs (16.2% faster)

def test_basic_zero_input():
    """Test that zero input returns 'Unknown size'."""
    codeflash_output = size_str(0) # 749ns -> 725ns (3.31% faster)

def test_basic_none_input():
    """Test that None input returns 'Unknown size'."""
    codeflash_output = size_str(None) # 671ns -> 670ns (0.149% faster)

def test_edge_negative_number():
    """Test behavior with negative numbers (should return 'Unknown size' due to falsy check)."""
    # The function checks `if not size_in_bytes:` which is falsy for negative zero,
    # but negative non-zero numbers are truthy and will proceed to calculation
    codeflash_output = size_str(-1024); result = codeflash_output # 3.23μs -> 2.63μs (23.0% faster)

def test_edge_very_small_positive_value():
    """Test with the smallest positive value that won't be falsy."""
    codeflash_output = size_str(0.1) # 2.95μs -> 2.21μs (33.3% faster)

def test_edge_bytes_just_below_kib():
    """Test with value just below 1 KiB threshold."""
    codeflash_output = size_str(1023) # 3.15μs -> 2.54μs (24.3% faster)

def test_edge_bytes_just_above_kib():
    """Test with value just above 1 KiB threshold."""
    codeflash_output = size_str(1025); result = codeflash_output # 4.50μs -> 3.88μs (16.1% faster)

def test_edge_mib_just_below_boundary():
    """Test with value just below 1 MiB."""
    codeflash_output = size_str(2**20 - 1); result = codeflash_output # 4.46μs -> 3.78μs (17.9% faster)

def test_edge_mib_just_above_boundary():
    """Test with value just above 1 MiB."""
    codeflash_output = size_str(2**20 + 1); result = codeflash_output # 4.39μs -> 3.71μs (18.2% faster)

def test_edge_gib_just_below_boundary():
    """Test with value just below 1 GiB."""
    codeflash_output = size_str(2**30 - 1); result = codeflash_output # 4.34μs -> 3.82μs (13.7% faster)

def test_edge_gib_just_above_boundary():
    """Test with value just above 1 GiB."""
    codeflash_output = size_str(2**30 + 1); result = codeflash_output # 4.27μs -> 3.86μs (10.6% faster)

def test_edge_tib_just_below_boundary():
    """Test with value just below 1 TiB."""
    codeflash_output = size_str(2**40 - 1); result = codeflash_output # 4.30μs -> 3.83μs (12.4% faster)

def test_edge_tib_just_above_boundary():
    """Test with value just above 1 TiB."""
    codeflash_output = size_str(2**40 + 1); result = codeflash_output # 4.03μs -> 3.73μs (7.91% faster)

def test_edge_pib_just_below_boundary():
    """Test with value just below 1 PiB."""
    codeflash_output = size_str(2**50 - 1); result = codeflash_output # 4.12μs -> 3.75μs (9.83% faster)

def test_edge_pib_just_above_boundary():
    """Test with value just above 1 PiB."""
    codeflash_output = size_str(2**50 + 1); result = codeflash_output # 3.91μs -> 3.56μs (9.77% faster)

def test_edge_very_large_pib_value():
    """Test with a very large value beyond 1000 PiB."""
    codeflash_output = size_str(1000 * 2**50); result = codeflash_output # 3.96μs -> 3.58μs (10.6% faster)

def test_edge_precision_two_decimal_places():
    """Test that output always has exactly 2 decimal places for non-byte values."""
    codeflash_output = size_str(int(1.234 * 2**20)); result = codeflash_output # 4.42μs -> 3.73μs (18.7% faster)

def test_edge_string_size_format():
    """Test that the output string has the correct format (number space unit)."""
    codeflash_output = size_str(2**20); result = codeflash_output # 4.38μs -> 3.73μs (17.2% faster)
    parts = result.split()

def test_edge_bytes_format_no_space_in_unit():
    """Test that bytes output doesn't have extra formatting."""
    codeflash_output = size_str(100); result = codeflash_output # 3.15μs -> 2.50μs (25.9% faster)

def test_edge_empty_string_input():
    """Test that empty string is treated as falsy and returns 'Unknown size'."""
    codeflash_output = size_str("") # 736ns -> 682ns (7.92% faster)

def test_edge_false_boolean_input():
    """Test that False is treated as falsy and returns 'Unknown size'."""
    codeflash_output = size_str(False) # 680ns -> 678ns (0.295% faster)

def test_edge_exactly_999_bytes():
    """Test edge case of 999 bytes (just under 1 KiB)."""
    codeflash_output = size_str(999) # 3.16μs -> 2.51μs (26.0% faster)

def test_edge_exactly_1000_bytes():
    """Test that 1000 bytes (not a power of 2) formats as bytes."""
    codeflash_output = size_str(1000) # 3.22μs -> 2.60μs (23.9% faster)

def test_edge_fractional_kilobyte():
    """Test fractional KiB value that's not a round number."""
    codeflash_output = size_str(int(1.456 * 2**10)); result = codeflash_output # 4.53μs -> 3.81μs (18.9% faster)

def test_edge_very_small_kilobyte():
    """Test value just above 1 KiB but less than 1.01 KiB."""
    codeflash_output = size_str(int(1.001 * 2**10)); result = codeflash_output # 4.46μs -> 3.78μs (18.0% faster)

def test_edge_exact_power_of_two_bytes():
    """Test various exact powers of two."""
    codeflash_output = size_str(2**10) # 4.38μs -> 3.82μs (14.6% faster)
    codeflash_output = size_str(2**11) # 1.41μs -> 1.06μs (32.3% faster)
    codeflash_output = size_str(2**20) # 1.08μs -> 820ns (31.6% faster)

def test_large_scale_progressively_larger_sizes():
    """Test a range of progressively larger sizes to ensure consistent formatting."""
    # Test sizes in different units
    test_cases = [
        (100, "100 bytes"),
        (2**10, "1.00 KiB"),
        (2**15, "32.00 KiB"),
        (2**20, "1.00 MiB"),
        (2**25, "32.00 MiB"),
        (2**30, "1.00 GiB"),
        (2**35, "32.00 GiB"),
        (2**40, "1.00 TiB"),
        (2**45, "32.00 TiB"),
        (2**50, "1.00 PiB"),
    ]
    for size, expected in test_cases:
        codeflash_output = size_str(size) # 13.3μs -> 11.3μs (17.3% faster)

def test_large_scale_many_byte_values():
    """Test 100 different byte values to ensure consistent behavior."""
    # Generate byte values from 1 to 1000
    for i in range(1, 1001):
        codeflash_output = size_str(i); result = codeflash_output # 835μs -> 649μs (28.7% faster)

def test_large_scale_many_kib_values():
    """Test 100 different KiB values to ensure consistent behavior."""
    # Generate KiB values from 1 to 1000
    for i in range(1, 101):
        size = i * 2**10
        codeflash_output = size_str(size); result = codeflash_output # 100μs -> 84.0μs (20.2% faster)

def test_large_scale_many_mib_values():
    """Test 100 different MiB values to ensure consistent behavior."""
    # Generate MiB values from 1 to 100
    for i in range(1, 101):
        size = i * 2**20
        codeflash_output = size_str(size); result = codeflash_output # 95.4μs -> 79.4μs (20.1% faster)

def test_large_scale_many_gib_values():
    """Test many GiB values to ensure consistent behavior."""
    # Generate GiB values from 1 to 100
    for i in range(1, 101):
        size = i * 2**30
        codeflash_output = size_str(size); result = codeflash_output # 91.6μs -> 75.1μs (22.0% faster)

def test_large_scale_many_tib_values():
    """Test many TiB values to ensure consistent behavior."""
    # Generate TiB values from 1 to 50
    for i in range(1, 51):
        size = i * 2**40
        codeflash_output = size_str(size); result = codeflash_output # 43.7μs -> 36.7μs (19.0% faster)

def test_large_scale_many_pib_values():
    """Test many PiB values to ensure consistent behavior."""
    # Generate PiB values from 1 to 50
    for i in range(1, 51):
        size = i * 2**50
        codeflash_output = size_str(size); result = codeflash_output # 39.6μs -> 34.4μs (15.2% faster)

def test_large_scale_fractional_sizes_across_all_units():
    """Test fractional sizes across all units."""
    test_cases = [
        (int(1.5 * 2**10), "1.50 KiB"),
        (int(2.7 * 2**10), "2.70 KiB"),
        (int(1.5 * 2**20), "1.50 MiB"),
        (int(2.7 * 2**20), "2.70 MiB"),
        (int(1.5 * 2**30), "1.50 GiB"),
        (int(2.7 * 2**30), "2.70 GiB"),
        (int(1.5 * 2**40), "1.50 TiB"),
        (int(2.7 * 2**40), "2.70 TiB"),
        (int(1.5 * 2**50), "1.50 PiB"),
        (int(2.7 * 2**50), "2.70 PiB"),
    ]
    for size, expected in test_cases:
        codeflash_output = size_str(size) # 12.9μs -> 11.0μs (16.6% faster)

def test_large_scale_random_sizes_format_validity():
    """Test that 200 random size values all produce valid formatted strings."""
    test_sizes = [
        1, 10, 100, 512, 1023, 1024, 2000, 5000,
        2**10, 2**11, 2**15, 2**19, 2**20, 2**21, 2**25, 2**29,
        2**30, 2**31, 2**35, 2**39, 2**40, 2**41, 2**45, 2**49,
        2**50, 2**51, 5 * 2**50, 10 * 2**50, 100 * 2**50, 500 * 2**50,
        int(1.1 * 2**10), int(1.5 * 2**20), int(2.3 * 2**30),
        int(3.7 * 2**40), int(4.2 * 2**50),
        999, 1025, (2**20 - 1), (2**20 + 1), (2**30 - 1), (2**30 + 1),
        (2**40 - 1), (2**40 + 1), (2**50 - 1), (2**50 + 1),
    ]
    for size in test_sizes:
        codeflash_output = size_str(size); result = codeflash_output # 44.7μs -> 38.1μs (17.4% faster)
        # Verify result contains a valid unit or "bytes"
        valid_units = ["bytes", "KiB", "MiB", "GiB", "TiB", "PiB"]
        # Verify result has the format "number unit"
        if "bytes" in result:
            parts = result.split()
        else:
            parts = result.split()
            try:
                float(parts[0])
            except ValueError:
                pytest.fail(f"Invalid number format in '{result}'")

def test_large_scale_monotonic_increase():
    """Test that size values monotonically increase in formatted size (roughly)."""
    sizes = [100, 1000, 10000, 100000, 1000000, 10000000, 100000000, 1000000000,
             10000000000, 100000000000, 1000000000000, 10000000000000]
    results = [size_str(s) for s in sizes]
    
    # All results should be valid strings with proper formatting
    for result in results:
        valid_units = ["bytes", "KiB", "MiB", "GiB", "TiB", "PiB"]

def test_large_scale_extreme_large_value():
    """Test with an extremely large value near the maximum representable size."""
    large_size = 999 * 2**50
    codeflash_output = size_str(large_size); result = codeflash_output # 3.98μs -> 3.63μs (9.60% faster)

def test_large_scale_consistency_multiple_calls():
    """Test that multiple calls to the function with the same input produce identical output."""
    test_size = int(3.7 * 2**35)
    results = [size_str(test_size) for _ in range(100)]

def test_large_scale_output_length_reasonable():
    """Test that output strings have reasonable lengths (not excessively long)."""
    test_cases = [100, 2**10, 2**20, 2**30, 2**40, 2**50, 999 * 2**50]
    for size in test_cases:
        codeflash_output = size_str(size); result = codeflash_output # 10.3μs -> 9.00μs (14.9% faster)

def test_large_scale_special_values_consistency():
    """Test that None, 0, and False all consistently return 'Unknown size'."""
    falsy_values = [None, 0, False, ""]
    for value in falsy_values:
        codeflash_output = size_str(value) # 1.45μs -> 1.32μs (9.91% faster)

def test_large_scale_decimal_precision_consistency():
    """Test that all non-byte outputs maintain consistent 2-decimal precision."""
    test_sizes = [
        int(1.234 * 2**10), int(5.678 * 2**20), int(9.999 * 2**30),
        int(1.001 * 2**40), int(100.999 * 2**50),
    ]
    for size in test_sizes:
        codeflash_output = size_str(size); result = codeflash_output # 9.02μs -> 7.95μs (13.4% faster)
        # For non-byte outputs, verify decimal format
        if "bytes" not in result:
            parts = result.split()
            number_str = parts[0]
            decimal_part = number_str.split(".")[1]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-size_str-mlcgmqgx and push.