While Python provides built-in iterators for sequences and collections, writing your own Custom Iterators allows you to stream complex data structures, generate custom mathematical sequences, and traverse domain models cleanly using Python's native for loop syntax.
1. Creating a Basic Custom Iterator
To create a custom iterator class, implement two essential methods:
1
__iter__(self): Returns the iterator instance (typically return self).
2
__next__(self): Calculates and returns the next element, or raises StopIteration when the sequence is completed.
Example 1: A Custom Countdown Iterator
1classCountdown:
2"""Iterates backwards from a given integer down to 1."""
3
4def__init__(self, start:int):
5 self.current = start
6
7def__iter__(self):
8# An iterator returns itself
9return self
10
11def__next__(self):
12if self.current <=0:
13raise StopIteration
14
15 value = self.current
16 self.current -=1
17return value
18
19# Using in a standard for loop
20for num in Countdown(5):
21print(num, end=" ")
22# Output: 5 4 3 2 1
23print()
2. Re-creating the range() Function
Here is how Python's built-in range() can be implemented from first principles:
1classStepRange:
2"""Generates integers from start up to (but not including) stop by step."""
Custom iterators can maintain complex state across iterations without keeping every generated number in memory:
1classFibonacci:
2"""Generates the first N numbers in the Fibonacci sequence."""
3
4def__init__(self, limit:int):
5 self.limit = limit
6 self.count =0
7 self.a, self.b =0,1
8
9def__iter__(self):
10return self
11
12def__next__(self):
13if self.count >= self.limit:
14raise StopIteration
15
16 val = self.a
17 self.a, self.b = self.b, self.a + self.b
18 self.count +=1
19return val
20
21fib = Fibonacci(8)
22print(list(fib))# [0, 1, 1, 2, 3, 5, 8, 13]
4. Re-usable Iterables: Separating Iterable from Iterator
When a class combines __iter__ and __next__ in one object, it is consumed after a single pass and cannot be re-iterated without re-instantiation.
To build an iterable collection that can be looped over multiple times (like a list), separate the Collection (Iterable) from the Iterator:
1classBookIterator:
2"""Maintains state for a single iteration pass."""
3def__init__(self, books):
4 self._books = books
5 self._index =0
6
7def__iter__(self):
8return self
9
10def__next__(self):
11if self._index >=len(self._books):
12raise StopIteration
13 book = self._books[self._index]
14 self._index +=1
15return book
16
17
18classBookshelf:
19"""Reusable Iterable collection."""
20def__init__(self):
21 self.books =[]
22
23defadd_book(self, title:str):
24 self.books.append(title)
25
26def__iter__(self):
27# Spawns a BRAND NEW iterator instance every time a loop starts!
28return BookIterator(self.books)
29
30shelf = Bookshelf()
31shelf.add_book("Clean Code")
32shelf.add_book("Fluent Python")
33shelf.add_book("Refactoring")
34
35# First loop works
36for b in shelf:
37print(f"Pass 1: {b}")
38
39# Second loop ALSO works because __iter__() returned a fresh iterator!
40for b in shelf:
41print(f"Pass 2: {b}")
5. Common Pitfalls
Forgetting StopIteration: If your __next__ method never raises StopIteration, loops over the iterator become infinite loops.
Forgetting return self: An iterator's __iter__ method must return an iterator object, typically self.
Modifying the underlying collection during iteration: Adding or removing elements from a list while iterating over it leads to skipped elements or index bugs.
⚠️ Common Intermediate Mistakes & Gotchas
1. Forgetting to Return self from __iter__()
An iterator must implement both __iter__() and __next__(). __iter__() should return self so the iterator can be used directly in for loops.
2. Forgetting to Raise StopIteration
If __next__() does not raise StopIteration upon completion, a for loop over the iterator will become an infinite loop.
💻 Try It Yourself: Step Range Iterator
Scenario
Build a custom float range iterator FloatRange(start, stop, step) that supports floating-point increments without floating-point drift.