Python Zero to Advanced #9: Iterators, Generators, Decorators & Context Managers

This part covers four features that separate "Python that works" from idiomatic, professional Python. They power frameworks you probably use: Flask and FastAPI routes are decorators, database sessions are context managers, and streaming large files relies on generators. Once you understand them, a lot of library "magic" stops being magic.

The iterator protocol

An iterable is anything you can loop over (list, str, dict, file). Calling iter() on it returns an iterator, an object with a __next__ method that returns the next value or raises StopIteration when exhausted. A for loop does exactly this behind the scenes.

colors = ["red", "green", "blue"]
it = iter(colors)
print(next(it))
print(next(it))
print(next(it))
try:
    next(it)
except StopIteration:
    print("exhausted")

class Countdown:
    def __init__(self, start):
        self.current = start

    def __iter__(self):
        return self

    def __next__(self):
        if self.current <= 0:
            raise StopIteration
        value = self.current
        self.current -= 1
        return value

print(list(Countdown(5)))

Output:

red
green
blue
exhausted
[5, 4, 3, 2, 1]

Note that an iterator can only be consumed once. After list(Countdown(5)) that particular object is empty.

Generators and yield

Writing iterator classes is verbose. A generator function uses yield instead of return; each call to next() runs the function until the next yield, then pauses and keeps its local state. Values are produced lazily, one at a time, so a generator can represent huge or even infinite sequences using almost no memory.

def countdown(start):
    print("starting")
    while start > 0:
        yield start
        start -= 1
    print("finished")

gen = countdown(3)
print(type(gen).__name__)
print(next(gen))
for n in gen:
    print(n)

def fibonacci():
    a, b = 0, 1
    while True:           # infinite, but lazy
        yield a
        a, b = b, a + b

from itertools import islice
print(list(islice(fibonacci(), 10)))

Output:

generator
starting
3
2
1
finished
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

Generator expressions look like list comprehensions with parentheses. They are perfect when you only need to iterate once, for example inside sum():

import sys

squares_list = [n * n for n in range(1_000_000)]
squares_gen = (n * n for n in range(1_000_000))

print(sum(n * n for n in range(1_000_000)))
print(sys.getsizeof(squares_list) > 1_000_000)
print(sys.getsizeof(squares_gen) < 1_000)

Output:

333332833333500000
True
True

Generator pipelines

Generators compose nicely into pipelines where each step processes one item at a time. This is how you can process a multi-gigabyte log file line by line:

lines = [
    "INFO user budi logged in",
    "ERROR payment timeout",
    "INFO user siti logged in",
    "ERROR db connection lost",
]

def read_lines(source):
    for line in source:          # could be an open file
        yield line.strip()

def only_errors(rows):
    return (r for r in rows if r.startswith("ERROR"))

def messages(rows):
    for r in rows:
        yield r.split(" ", 1)[1]

for msg in messages(only_errors(read_lines(lines))):
    print(msg)

Output:

payment timeout
db connection lost

Decorators

A decorator is a function that takes a function and returns a new function that usually wraps the original with extra behaviour. The @decorator syntax above a def is just shorthand for func = decorator(func). Always use functools.wraps so the wrapped function keeps its name and docstring.

from functools import wraps

def log_calls(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f"calling {func.__name__}{args} {kwargs}")
        result = func(*args, **kwargs)
        print(f"{func.__name__} returned {result!r}")
        return result
    return wrapper

@log_calls
def add(a, b, round_to=None):
    """Add two numbers."""
    total = a + b
    return round(total, round_to) if round_to is not None else total

add(2, 3)
add(1.234, 2.1, round_to=1)
print(add.__name__, "-", add.__doc__)

Output:

calling add(2, 3) {}
add returned 5
calling add(1.234, 2.1) {'round_to': 1}
add returned 3.3
add - Add two numbers.

Decorators with arguments

To pass arguments to a decorator you add one more layer: a function that returns the actual decorator.

from functools import wraps

def retry(times):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except ConnectionError as e:
                    print(f"attempt {attempt} failed: {e}")
            raise ConnectionError(f"gave up after {times} attempts")
        return wrapper
    return decorator

calls = {"n": 0}

@retry(times=3)
def fetch_data():
    calls["n"] += 1
    if calls["n"] < 3:
        raise ConnectionError("timeout")
    return {"status": "ok"}

print(fetch_data())

Output:

attempt 1 failed: timeout
attempt 2 failed: timeout
{'status': 'ok'}

The standard library includes useful decorators too, such as functools.lru_cache for memoisation:

from functools import lru_cache

@lru_cache(maxsize=None)
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

print(fib(80))
print(fib.cache_info().hits)

Output:

23416728348467685
78

Context managers

You already used one: with open(...). A context manager runs setup code on entering the with block and guaranteed teardown code on exit, even if an exception happens. You can write one as a class with __enter__/__exit__, or more simply with contextlib.contextmanager and a generator.

from contextlib import contextmanager

class Transaction:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"BEGIN {self.name}")
        return self

    def __exit__(self, exc_type, exc, tb):
        if exc_type is None:
            print("COMMIT")
        else:
            print(f"ROLLBACK because {exc_type.__name__}: {exc}")
        return True          # True = suppress the exception

with Transaction("transfer") as tx:
    print("debit Budi")
    print("credit Siti")

with Transaction("refund"):
    print("debit Siti")
    raise ValueError("amount mismatch")

@contextmanager
def tag(name):
    print(f"<{name}>")
    try:
        yield
    finally:
        print(f"</{name}>")

with tag("section"):
    print("content")

Output:

BEGIN transfer
debit Budi
credit Siti
COMMIT
BEGIN refund
debit Siti
ROLLBACK because ValueError: amount mismatch
<section>
content
</section>

Returning True from __exit__ swallows the exception, so use it carefully; most context managers return None (falsy) and let errors propagate.

Common mistakes

• Iterating a generator twice. The second loop gets nothing. Recreate it or convert to a list if you need multiple passes.
• Forgetting @wraps: every decorated function is then named wrapper, which confuses logs and tools.
• Forgetting to return the result inside a wrapper, so the decorated function always returns None.
• Not using try/finally around yield in a @contextmanager, so cleanup is skipped on errors.

Exercise

Write a generator chunked(items, size) that yields lists of at most size items, useful for sending records to an API in batches.

def chunked(items, size):
    batch = []
    for item in items:
        batch.append(item)
        if len(batch) == size:
            yield batch
            batch = []
    if batch:
        yield batch

for batch in chunked(range(1, 11), 4):
    print(batch)

Output:

[1, 2, 3, 4]
[5, 6, 7, 8]
[9, 10]

Summary

• Iterators implement __iter__/__next__; for loops use them automatically.
• Generators (yield and generator expressions) produce values lazily and save memory.
• Decorators wrap functions with extra behaviour; use functools.wraps.
• Context managers guarantee setup and cleanup with with.

In the final part, Part 10, we cover type hints, asyncio, concurrency basics, and testing with pytest.