Custom Iterator Class
Overview
A custom iterator class is a class that implements both the `__iter__()` and `__next__()` dunder methods. This allows its instances to be looped over directly with a `for` loop. - `__iter__()`: Returns the iterator object itself (usually `self`). - `__next__()`: Computes and returns the next value in the sequence. Raises `StopIteration` when exhausted. Think of it like a Pez dispenser: the dispenser is the iterator (handling `__iter__()`), and each time you push the head back, you call `__next__()` to get a candy. Once empty, it raises a signal that there's no more candy. A common trap is assuming an iterator resets. Because iterators hold state, they are completely exhausted after a single pass: iterator = CountTwo() list(iterator) # [1, 2] list(iterator) # [] (Trap: already exhausted) This behavior is covered fully in generator-expressions where similar lazy evaluation patterns are used.
It allows you to create custom, memory-efficient data streams that compute or fetch values lazily, instead of loading a massive `list` into memory at once.
Where used: FastAPI, Django ORM, Data Pipelines
Why learn this
- Build memory-efficient paginators for APIs.
- Understand how Python's `for` loop actually works under the hood.
- Create streaming data processors.
Code walkthrough
class CountTwo:
def __init__(self):
self.count = 0
def __iter__(self):
return self
def __next__(self):
if self.count >= 2:
raise StopIteration
self.count += 1
return self.count
for n in CountTwo():
print(n)Focus: if self.count >= 2: raise StopIteration
Aha moment
class Exhaustible:
def __init__(self):
self.val = 2
def __iter__(self):
return self
def __next__(self):
if self.val == 0: raise StopIteration
self.val -= 1
return self.val
iterator = Exhaustible()
print(list(iterator))
print(list(iterator))Prediction: What will the two `print()` statements output?
Common guess: `[1, 0]` followed by `[1, 0]`
Iterators maintain their internal state. Once `self.val` reaches `0` and raises `StopIteration`, the iterator is exhausted. A second attempt to iterate over it does nothing and returns an empty `list`.
Common mistakes
- Forgetting to raise `StopIteration`: If `__next__()` does not raise `StopIteration`, a `for` loop over the object will run infinitely.
- Not returning `self` in `__iter__()`: If `__iter__()` does not return `self` on an iterator, passing the iterator itself into a `for` loop or `iter()` will fail with a `TypeError`.
Glossary
- exhausted
- When an iterator has gone through all its items and has no more left to provide, e.g. `next(it)` raises `StopIteration`.
- lazily
- Computing or fetching values only exactly when they are needed, rather than all at once in advance, e.g. `line = next(f)`.
Recall questions
- What two dunder methods must a class implement to be a complete custom iterator?
- What exception must `__next__()` raise to signal that the iteration is finished?
- What should the `__iter__()` method typically return when defined on an iterator class?
Understanding checks
What is the output of this code?
`[1, 2]` followed by `[]`
Iterators hold state. Once the first `list(c)` exhausts the iterator by raising `StopIteration`, its internal state (`self.num`) remains at `2`. The second `list(c)` immediately hits the `StopIteration` condition.
This custom iterator is supposed to yield numbers `1`, `2`, `3` but it causes a crash. Why?
The `__iter__()` method is missing `return self`.
A standard `for` loop first calls `iter()` on the object, which invokes `__iter__()`. The `iter()` function expects an iterator object (one with a `__next__()` method) to be returned. Because `__iter__()` implicitly returns `None`, `iter()` immediately rejects it and raises a `TypeError`.
Practice tasks
Limit Iterator Output
The `CycleList` iterator infinitely loops over a `list`. Modify it to accept a `max_yields` argument in its constructor. Update `__next__()` to raise `StopIteration` once it has yielded that many items.
Challenge
Pagination Iterator
Create an iterator class `Paginator` that takes a `list` of `items` and a `page_size`. Each iteration should return the next slice of the `list` of size `page_size`. When no items remain, raise `StopIteration`.
Continue learning
Previous: __iter__ and __next__
Previous: Everything is an object in Python