Iteration

In this chapter, we'll explore iteration in Python - one of the most powerful features of the language. We'll learn:

• for loops with ranges and collections
• while loops for condition-based iteration
• List comprehensions for elegant transformations
• The sum() function
• Benchmarking with the timeit module

First example: Sum

Let's write a function that sums a list of numbers.

Write the test first

Create test_iteration.py:

from iteration import sum_numbers


def test_sum_numbers():
    numbers = [1, 2, 3, 4, 5]
    got = sum_numbers(numbers)
    want = 15
    assert got == want

Try to run it

pytest test_iteration.py

ModuleNotFoundError: No module named 'iteration'

Write the minimal code

Create iteration.py:

def sum_numbers(numbers):
    pass

Run the test:

AssertionError: assert None == 15

Make it pass

def sum_numbers(numbers):
    result = 0
    for number in numbers:
        result += number
    return result

The for loop iterates over each element in the list. Python's for loop is actually a "for each" loop - it gives you each element directly, not an index.

Refactor with built-in sum

Python has a built-in sum() function:

def sum_numbers(numbers):
    return sum(numbers)

Run the tests to make sure they still pass!

Write more tests

Let's add edge cases:

from iteration import sum_numbers


def test_sum_numbers():
    numbers = [1, 2, 3, 4, 5]
    got = sum_numbers(numbers)
    want = 15
    assert got == want


def test_sum_empty_list():
    got = sum_numbers([])
    want = 0
    assert got == want


def test_sum_negative_numbers():
    got = sum_numbers([-1, -2, -3])
    want = -6
    assert got == want


def test_sum_single_element():
    got = sum_numbers([42])
    want = 42
    assert got == want

All tests should pass!

Repeat: The repeat function

Let's write a function that repeats a character a specified number of times.

Write the test first

def test_repeat():
    got = repeat("a", 5)
    want = "aaaaa"
    assert got == want

Make it compile

def repeat(character, count):
    pass

Make it pass using a for loop

def repeat(character, count):
    result = ""
    for _ in range(count):
        result += character
    return result

New concepts:

• range(count) generates numbers from 0 to count-1
• _ is a convention for a variable we don't need to use
• Strings can be concatenated with +

Refactor

Python has a simpler way - string multiplication:

def repeat(character, count):
    return character * count

Run the tests to verify!

The range function

range() is essential for iteration:

range(5)        # 0, 1, 2, 3, 4
range(1, 5)     # 1, 2, 3, 4
range(0, 10, 2) # 0, 2, 4, 6, 8 (step of 2)
range(5, 0, -1) # 5, 4, 3, 2, 1 (counting down)

while loops

while loops repeat as long as a condition is true:

def countdown(n):
    """Return a countdown from n to 1."""
    result = []
    while n > 0:
        result.append(n)
        n -= 1
    return result

Test it:

def test_countdown():
    got = countdown(5)
    want = [5, 4, 3, 2, 1]
    assert got == want


def test_countdown_zero():
    got = countdown(0)
    want = []
    assert got == want

List comprehensions

List comprehensions are a concise way to create lists:

# Traditional for loop
squares = []
for x in range(5):
    squares.append(x ** 2)
# Result: [0, 1, 4, 9, 16]

# List comprehension
squares = [x ** 2 for x in range(5)]
# Result: [0, 1, 4, 9, 16]

Filtering with comprehensions

You can add conditions to filter elements:

# Get even numbers
evens = [x for x in range(10) if x % 2 == 0]
# Result: [0, 2, 4, 6, 8]

# Get positive numbers from a list
positives = [x for x in numbers if x > 0]

Testing list comprehensions

Let's write a function that filters to only positive numbers:

def test_positive_only():
    numbers = [-3, -1, 0, 1, 3, 5]
    got = positive_only(numbers)
    want = [1, 3, 5]
    assert got == want

Implementation:

def positive_only(numbers):
    return [x for x in numbers if x > 0]

Fibonacci sequence

Let's implement the Fibonacci sequence - a classic programming exercise.

Write the test first

def test_fibonacci():
    test_cases = [
        (0, 0),
        (1, 1),
        (2, 1),
        (3, 2),
        (4, 3),
        (5, 5),
        (10, 55),
    ]
    
    for n, expected in test_cases:
        got = fibonacci(n)
        assert got == expected, f"fibonacci({n}) = {got}, want {expected}"

Make it pass

def fibonacci(n):
    """Return the nth Fibonacci number."""
    if n <= 1:
        return n
    
    a, b = 0, 1
    for _ in range(2, n + 1):
        a, b = b, a + b
    return b

New concepts:

• a, b = 0, 1 is tuple unpacking - assigns multiple values at once
• a, b = b, a + b swaps and updates values simultaneously

Enumerate

When you need both the index and value during iteration:

fruits = ["apple", "banana", "cherry"]

# Without enumerate
for i in range(len(fruits)):
    print(f"{i}: {fruits[i]}")

# With enumerate (preferred)
for i, fruit in enumerate(fruits):
    print(f"{i}: {fruit}")

Zip

To iterate over multiple lists in parallel:

names = ["Alice", "Bob", "Charlie"]
ages = [25, 30, 35]

for name, age in zip(names, ages):
    print(f"{name} is {age}")

Benchmarking with timeit

When you want to compare performance:

import timeit

# Time a string concatenation loop
def concat_loop():
    result = ""
    for i in range(1000):
        result += str(i)
    return result

# Time a join approach
def concat_join():
    return "".join(str(i) for i in range(1000))

# Benchmark
loop_time = timeit.timeit(concat_loop, number=1000)
join_time = timeit.timeit(concat_join, number=1000)

print(f"Loop: {loop_time:.4f}s")
print(f"Join: {join_time:.4f}s")

The join approach is typically faster because strings are immutable in Python, so concatenation creates new strings each time.

Break and continue

Control loop flow with break and continue:

# break - exit the loop entirely
for i in range(10):
    if i == 5:
        break
    print(i)  # Prints 0, 1, 2, 3, 4

# continue - skip to the next iteration
for i in range(5):
    if i == 2:
        continue
    print(i)  # Prints 0, 1, 3, 4

The else clause on loops

Python loops can have an else clause that runs when the loop completes without break:

def find_first_even(numbers):
    for num in numbers:
        if num % 2 == 0:
            return num
    else:
        return None  # No even number found

# Test
def test_find_first_even():
    assert find_first_even([1, 3, 4, 5]) == 4
    assert find_first_even([1, 3, 5]) is None

Wrapping up

We've covered:

• for loops - Iterate over collections with for item in collection
• while loops - Repeat while a condition is true
• range() - Generate sequences of numbers
• List comprehensions - Concise syntax for creating lists: [x for x in items if condition]
• enumerate() - Get index and value together
• zip() - Iterate over multiple lists in parallel
• break and continue - Control loop flow
• timeit - Benchmark code performance

The TDD approach

Throughout this chapter, we:

1. Wrote tests first to define expected behavior
2. Made tests pass with the simplest code
3. Refactored to more Pythonic solutions
4. Kept tests passing throughout

This cycle of test → implement → refactor is the heart of TDD.