⚡️ Speed up function _get_total_scope_state by 22%

[codeflash-ai]

📄 22% (0.22x) speedup for _get_total_scope_state in src/uberjob/progress/_html_progress_observer.py

⏱️ Runtime : 1.20 milliseconds → 986 microseconds (best of 250 runs)

📝 Explanation and details

The optimization replaces five separate sum() operations with generator expressions with a single loop that accumulates all fields simultaneously.

Key optimization: Instead of iterating over scope_states five times (once per field), the optimized version iterates only once, accumulating all five fields (completed, failed, running, total, weighted_elapsed) in each iteration.

Why this is faster:

• Reduced iterations: Eliminates 4 out of 5 passes through the collection, significantly reducing overhead as collection size grows
• Better cache locality: Accessing all fields of each ScopeState object while it's already in CPU cache is more efficient than repeated separate iterations
• Lower function call overhead: Eliminates the overhead of creating and executing five separate generator expressions

Performance characteristics: The optimization shows greater speedup with smaller collections (71-77% faster for single elements) and moderate improvements for larger collections (13-35% faster for 1000+ elements). This pattern suggests the optimization primarily reduces fixed overhead per iteration rather than per-element costs.

Impact on workloads: Based on the function reference, _get_total_scope_state is called from _render_section when there are multiple scopes to display progress totals. Since progress rendering typically happens frequently during job execution, this 21% speedup will meaningfully improve responsiveness of progress displays, especially when aggregating multiple scope states.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from collections import namedtuple

# imports
import pytest  # used for our unit tests
from uberjob.progress._html_progress_observer import _get_total_scope_state


# function to test
class ScopeState:
    """Simple data structure for scope state."""

    def __init__(self, completed, failed, running, total, weighted_elapsed):
        self.completed = completed
        self.failed = failed
        self.running = running
        self.total = total
        self.weighted_elapsed = weighted_elapsed

    def __eq__(self, other):
        if not isinstance(other, ScopeState):
            return False
        return (
            self.completed == other.completed
            and self.failed == other.failed
            and self.running == other.running
            and self.total == other.total
            and self.weighted_elapsed == other.weighted_elapsed
        )

    def __repr__(self):
        return (
            f"ScopeState(completed={self.completed}, failed={self.failed}, "
            f"running={self.running}, total={self.total}, weighted_elapsed={self.weighted_elapsed})"
        )


from uberjob.progress._html_progress_observer import _get_total_scope_state

# unit tests

# ------------------ Basic Test Cases ------------------


def test_empty_scope_states():
    """Test with an empty list of scope states."""
    codeflash_output = _get_total_scope_state([])
    result = codeflash_output  # 2.65μs -> 1.51μs (75.9% faster)


def test_single_scope_state():
    """Test with a single ScopeState object."""
    s = ScopeState(completed=1, failed=2, running=3, total=4, weighted_elapsed=5)
    codeflash_output = _get_total_scope_state([s])
    result = codeflash_output  # 2.69μs -> 1.57μs (71.5% faster)


def test_multiple_scope_states_basic():
    """Test with multiple ScopeState objects with positive integers."""
    s1 = ScopeState(1, 2, 3, 4, 5)
    s2 = ScopeState(10, 20, 30, 40, 50)
    s3 = ScopeState(100, 200, 300, 400, 500)
    codeflash_output = _get_total_scope_state([s1, s2, s3])
    result = codeflash_output  # 3.13μs -> 2.13μs (47.1% faster)
    expected = ScopeState(
        1 + 10 + 100, 2 + 20 + 200, 3 + 30 + 300, 4 + 40 + 400, 5 + 50 + 500
    )


def test_multiple_scope_states_with_zeros():
    """Test with multiple ScopeState objects where some have zeros."""
    s1 = ScopeState(0, 0, 0, 0, 0)
    s2 = ScopeState(1, 2, 3, 4, 5)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.81μs -> 1.85μs (51.7% faster)


# ------------------ Edge Test Cases ------------------


def test_scope_states_with_negative_values():
    """Test with negative values in ScopeState fields."""
    s1 = ScopeState(-1, -2, -3, -4, -5)
    s2 = ScopeState(5, 4, 3, 2, 1)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.71μs -> 1.77μs (53.4% faster)
    expected = ScopeState(4, 2, 0, -2, -4)


def test_scope_states_with_large_numbers():
    """Test with very large numbers to check for integer overflow."""
    large = 10**18
    s1 = ScopeState(large, large, large, large, large)
    s2 = ScopeState(large, large, large, large, large)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.85μs -> 2.03μs (39.9% faster)
    expected = ScopeState(large * 2, large * 2, large * 2, large * 2, large * 2)


def test_scope_states_with_floats():
    """Test with float values in weighted_elapsed (and possibly other fields)."""
    s1 = ScopeState(1, 2, 3, 4, 1.5)
    s2 = ScopeState(10, 20, 30, 40, 2.5)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 3.18μs -> 2.07μs (53.8% faster)
    expected = ScopeState(11, 22, 33, 44, 4.0)


def test_scope_states_with_mixed_types():
    """Test with a mix of integer and float values."""
    s1 = ScopeState(1, 2, 3, 4, 1.0)
    s2 = ScopeState(10, 20, 30, 40, 2)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.93μs -> 2.04μs (43.8% faster)
    expected = ScopeState(11, 22, 33, 44, 3.0)


def test_scope_states_with_all_zero_fields():
    """Test with all fields set to zero."""
    s1 = ScopeState(0, 0, 0, 0, 0)
    s2 = ScopeState(0, 0, 0, 0, 0)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.72μs -> 1.78μs (52.8% faster)


def test_scope_states_with_one_field_nonzero():
    """Test with only one field nonzero in one ScopeState."""
    s1 = ScopeState(0, 0, 0, 0, 0)
    s2 = ScopeState(0, 0, 1, 0, 0)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.67μs -> 1.77μs (50.7% faster)


def test_scope_states_with_non_scope_state_object():
    """Test that non-ScopeState objects raise AttributeError."""

    class Dummy:
        pass

    with pytest.raises(AttributeError):
        _get_total_scope_state([Dummy()])  # 2.06μs -> 1.55μs (32.4% faster)


def test_scope_states_with_namedtuple():
    """Test with a namedtuple that matches the ScopeState interface."""
    MyScopeState = namedtuple(
        "MyScopeState", ["completed", "failed", "running", "total", "weighted_elapsed"]
    )
    s1 = MyScopeState(1, 2, 3, 4, 5)
    s2 = MyScopeState(10, 20, 30, 40, 50)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 3.85μs -> 2.40μs (60.1% faster)
    expected = ScopeState(11, 22, 33, 44, 55)


def test_scope_states_with_missing_field():
    """Test with an object missing one of the required fields."""

    class IncompleteScopeState:
        def __init__(self):
            self.completed = 1
            self.failed = 2
            self.running = 3
            self.total = 4
            # missing weighted_elapsed

    with pytest.raises(AttributeError):
        _get_total_scope_state(
            [IncompleteScopeState()]
        )  # 3.16μs -> 1.85μs (71.2% faster)


# ------------------ Large Scale Test Cases ------------------


def test_large_scale_scope_states_all_ones():
    """Test with a large number of ScopeState objects, all fields set to 1."""
    n = 1000
    states = [ScopeState(1, 1, 1, 1, 1) for _ in range(n)]
    codeflash_output = _get_total_scope_state(states)
    result = codeflash_output  # 151μs -> 121μs (25.3% faster)
    expected = ScopeState(n, n, n, n, n)


def test_large_scale_scope_states_varied():
    """Test with a large number of ScopeState objects with varied values."""
    n = 1000
    states = [ScopeState(i, i * 2, i * 3, i * 4, float(i * 5)) for i in range(n)]
    codeflash_output = _get_total_scope_state(states)
    result = codeflash_output  # 140μs -> 124μs (13.1% faster)
    completed_sum = sum(i for i in range(n))
    failed_sum = sum(i * 2 for i in range(n))
    running_sum = sum(i * 3 for i in range(n))
    total_sum = sum(i * 4 for i in range(n))
    weighted_elapsed_sum = sum(i * 5 for i in range(n))
    expected = ScopeState(
        completed_sum, failed_sum, running_sum, total_sum, float(weighted_elapsed_sum)
    )


def test_large_scale_scope_states_large_values():
    """Test with a large number of ScopeState objects with large values."""
    n = 1000
    large = 10**12
    states = [ScopeState(large, large, large, large, large) for _ in range(n)]
    codeflash_output = _get_total_scope_state(states)
    result = codeflash_output  # 143μs -> 140μs (1.67% faster)
    expected = ScopeState(large * n, large * n, large * n, large * n, large * n)


def test_large_scale_scope_states_mixed_zero_and_nonzero():
    """Test with half ScopeState objects zero, half nonzero."""
    n = 1000
    states = [ScopeState(0, 0, 0, 0, 0) for _ in range(n // 2)] + [
        ScopeState(1, 2, 3, 4, 5) for _ in range(n // 2)
    ]
    codeflash_output = _get_total_scope_state(states)
    result = codeflash_output  # 140μs -> 113μs (23.2% faster)
    expected = ScopeState(n // 2 * 1, n // 2 * 2, n // 2 * 3, n // 2 * 4, n // 2 * 5)


def test_large_scale_scope_states_with_floats():
    """Test with a large number of ScopeState objects with float weighted_elapsed."""
    n = 1000
    states = [ScopeState(1, 2, 3, 4, float(i) / 10) for i in range(n)]
    codeflash_output = _get_total_scope_state(states)
    result = codeflash_output  # 136μs -> 121μs (12.7% faster)
    completed_sum = n * 1
    failed_sum = n * 2
    running_sum = n * 3
    total_sum = n * 4
    weighted_elapsed_sum = sum(float(i) / 10 for i in range(n))
    expected = ScopeState(
        completed_sum, failed_sum, running_sum, total_sum, weighted_elapsed_sum
    )


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from collections import namedtuple

# imports
import pytest
from uberjob.progress._html_progress_observer import _get_total_scope_state

# Define ScopeState as in the source (since we don't have the actual import)
ScopeState = namedtuple(
    "ScopeState", ["completed", "failed", "running", "total", "weighted_elapsed"]
)
from uberjob.progress._html_progress_observer import _get_total_scope_state

# unit tests

# -------- BASIC TEST CASES --------


def test_empty_scope_states():
    # Test with an empty list - all sums should be zero
    codeflash_output = _get_total_scope_state([])
    result = codeflash_output  # 2.49μs -> 1.41μs (77.0% faster)


def test_single_scope_state():
    # Test with a single ScopeState - should return the same values
    s = ScopeState(completed=1, failed=2, running=3, total=6, weighted_elapsed=7)
    codeflash_output = _get_total_scope_state([s])
    result = codeflash_output  # 2.70μs -> 1.92μs (40.7% faster)


def test_multiple_scope_states_basic():
    # Test with a few ScopeStates with small positive values
    s1 = ScopeState(1, 0, 2, 3, 4)
    s2 = ScopeState(2, 1, 0, 3, 1)
    s3 = ScopeState(0, 1, 2, 3, 5)
    codeflash_output = _get_total_scope_state([s1, s2, s3])
    result = codeflash_output  # 3.24μs -> 2.22μs (45.7% faster)


def test_scope_states_with_zeros():
    # Test when all ScopeStates have zero for all fields
    s1 = ScopeState(0, 0, 0, 0, 0)
    s2 = ScopeState(0, 0, 0, 0, 0)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.81μs -> 1.89μs (48.8% faster)


def test_scope_states_with_negative_values():
    # Test with negative values (should sum correctly)
    s1 = ScopeState(-1, -2, -3, -4, -5)
    s2 = ScopeState(1, 2, 3, 4, 5)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.78μs -> 1.90μs (45.9% faster)


def test_scope_states_with_mixed_signs():
    # Test with mixed positive and negative values
    s1 = ScopeState(5, -2, 3, -4, 10)
    s2 = ScopeState(-1, 2, -3, 4, -10)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.84μs -> 1.94μs (46.5% faster)


# -------- EDGE TEST CASES --------


def test_scope_states_with_large_integers():
    # Test with very large integer values
    big = 10**18
    s1 = ScopeState(big, big, big, big, big)
    s2 = ScopeState(big, big, big, big, big)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 2.88μs -> 2.14μs (34.4% faster)


def test_scope_states_with_min_max_integers():
    # Test with min and max integer values (Python ints are unbounded, but test extremes)
    min_int = -(2**63)
    max_int = 2**63 - 1
    s1 = ScopeState(min_int, min_int, min_int, min_int, min_int)
    s2 = ScopeState(max_int, max_int, max_int, max_int, max_int)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 3.15μs -> 2.41μs (30.7% faster)
    expected = ScopeState(
        min_int + max_int,
        min_int + max_int,
        min_int + max_int,
        min_int + max_int,
        min_int + max_int,
    )


def test_scope_states_with_floats():
    # Test with float values for weighted_elapsed
    s1 = ScopeState(1, 2, 3, 4, 1.5)
    s2 = ScopeState(5, 6, 7, 8, 2.5)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 3.19μs -> 2.16μs (47.5% faster)


def test_scope_states_with_only_weighted_elapsed_nonzero():
    # Only weighted_elapsed has nonzero values
    s1 = ScopeState(0, 0, 0, 0, 1.1)
    s2 = ScopeState(0, 0, 0, 0, 2.2)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 3.08μs -> 2.09μs (47.4% faster)


def test_scope_states_with_non_integer_types():
    # Test with float values in integer fields (should sum as floats)
    s1 = ScopeState(1.1, 2.2, 3.3, 4.4, 5.5)
    s2 = ScopeState(6.6, 7.7, 8.8, 9.9, 10.1)
    codeflash_output = _get_total_scope_state([s1, s2])
    result = codeflash_output  # 3.08μs -> 2.13μs (44.6% faster)


def test_scope_states_with_non_scope_state_objects():
    # Test that it raises AttributeError if input is not ScopeState-like
    with pytest.raises(AttributeError):
        _get_total_scope_state([1, 2, 3])  # 1.95μs -> 1.46μs (33.5% faster)


def test_scope_states_with_none():
    # Test that it raises AttributeError if input contains None
    s = ScopeState(1, 2, 3, 4, 5)
    with pytest.raises(AttributeError):
        _get_total_scope_state([s, None])  # 2.06μs -> 1.95μs (5.43% faster)


# -------- LARGE SCALE TEST CASES --------


def test_large_number_of_scope_states():
    # Test with a large number of ScopeStates (1000)
    n = 1000
    s = ScopeState(1, 2, 3, 4, 5)
    scope_states = [s] * n
    codeflash_output = _get_total_scope_state(scope_states)
    result = codeflash_output  # 161μs -> 121μs (32.7% faster)


def test_large_number_of_varied_scope_states():
    # Test with 1000 ScopeStates with increasing values
    n = 1000
    scope_states = [ScopeState(i, i + 1, i + 2, i + 3, i + 4) for i in range(n)]
    # Calculate expected sums
    completed_sum = sum(i for i in range(n))
    failed_sum = sum(i + 1 for i in range(n))
    running_sum = sum(i + 2 for i in range(n))
    total_sum = sum(i + 3 for i in range(n))
    weighted_elapsed_sum = sum(i + 4 for i in range(n))
    codeflash_output = _get_total_scope_state(scope_states)
    result = codeflash_output  # 167μs -> 123μs (35.4% faster)


def test_large_scope_states_with_large_numbers():
    # Test with 500 ScopeStates with large values
    n = 500
    val = 10**12
    scope_states = [ScopeState(val, val, val, val, val) for _ in range(n)]
    codeflash_output = _get_total_scope_state(scope_states)
    result = codeflash_output  # 85.2μs -> 68.0μs (25.3% faster)


# 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-_get_total_scope_state-mi6tfyqj and push.