PYTHON / ITERATORS, GENERATORS, AND COMPREHENSIONS
Building an iterator class
Write your own iterator classes with __iter__ and __next__, and choose between one-shot iterators and re-iterable containers.
What you will learn
- Implement __next__ to advance stored state and raise StopIteration when done
- Return self from __iter__ for a one-shot iterator, a new object for a reusable one
- Explain why list(obj) gives data the first time and [] the second
- Wrap another iterable with iter() in __init__ to build batching or filtering iterators
Understanding Building an iterator class
An iterator class is an object that owns a cursor. The cursor lives in instance attributes set up by __init__, and __next__ is the only method allowed to move it: it either computes the next value, mutates the state, and returns the value, or it raises StopIteration to say the cursor has run off the end. __iter__ exists so the object works with a for statement; when the object is its own cursor, __iter__ just returns self.
That self-returning design has a direct consequence: the instance is one-shot. A for loop calls iter() exactly once and then calls __next__ until StopIteration, so a partially consumed instance resumes where it stopped, and a fully consumed one produces nothing at all on the next loop. This is not a defect; open file objects and generators behave the same way, which is what makes them safe to pass around without secretly rewinding.
When you want the object to be looped over repeatedly, split the two roles: a container class whose __iter__ constructs and returns a brand new iterator object holding the cursor. This is exactly why a list is iterable but not an iterator; iter([1, 2]) hands back a fresh list_iterator each call, and that freshness is what lets zip(xs, xs) pair items and nested for loops over the same list work as expected.
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
c = Countdown(3)
print([n for n in c])
print([n for n in c])
print(iter(c) is c)An iterator holds mutable cursor state in __next__, so whether __iter__ returns self or a new object decides if the thing can be looped over more than once.
Worked examples
Re-iterable container with a separate iterator
Splitting the container from the cursor lets the same object be looped over many times, and independently.
class Repeat:
def __init__(self, values, times):
self.values = values
self.times = times
def __iter__(self):
return RepeatIterator(self.values, self.times)
class RepeatIterator:
def __init__(self, values, times):
self.values = values
self.remaining = times
self.index = 0
def __iter__(self):
return self
def __next__(self):
if self.remaining == 0:
raise StopIteration
value = self.values[self.index]
self.index += 1
if self.index == len(self.values):
self.index = 0
self.remaining -= 1
return value
r = Repeat(["a", "b"], 2)
print(list(r))
print(list(r))
print(list(zip(r, r)))Example explained
Line 1Repeat.__iter__ builds a new RepeatIterator each call, so the second list(r) starts from scratch.
Line 2RepeatIterator keeps two pieces of state, index within values and remaining passes, and only __next__ touches them.
Line 3RepeatIterator still defines __iter__ returning self, which is required for it to be usable in a for loop itself.
Line 4zip(r, r) calls iter(r) twice and gets two unrelated cursors, so the pairs line up instead of interleaving.
Watching the protocol calls
Prints inside the methods show that for calls __iter__ once and __next__ until StopIteration.
class Chars:
def __init__(self, text):
self.text = text
self.pos = 0
def __iter__(self):
print("__iter__ called")
return self
def __next__(self):
if self.pos >= len(self.text):
print("raising StopIteration")
raise StopIteration
ch = self.text[self.pos]
self.pos += 1
return ch
it = Chars("hi")
for ch in it:
print("got", ch)
print(next(it, "<empty>"))Example explained
Line 1"__iter__ called" appears once: the loop resolves the iterator a single time, before the first item.
Line 2The guard uses self.pos >= len(self.text) before reading, so the last valid character is returned and the stop comes on the following call.
Line 3StopIteration is raised twice, once to end the loop and once for the manual next(), proving the exhausted state is permanent.
Line 4next(it, "<empty>") catches that StopIteration and substitutes the default instead of propagating it.
An iterator that wraps another iterable
Holding iter(source) as state lets one __next__ call consume several items from the underlying iterable.
class Batched:
def __init__(self, iterable, size):
self.source = iter(iterable)
self.size = size
def __iter__(self):
return self
def __next__(self):
batch = []
for item in self.source:
batch.append(item)
if len(batch) == self.size:
return batch
if batch:
return batch
raise StopIteration
print(list(Batched(range(7), 3)))
b = Batched("abcd", 2)
print(next(b), next(b))
print(next(b, None))Example explained
Line 1iter(iterable) in __init__ stores a cursor over the source, so successive __next__ calls resume rather than restart.
Line 2The inner for loop stops as soon as the batch is full and returns, leaving the source positioned for the next call.
Line 3A short final batch is returned only if it is non-empty, which is how [6] appears without an extra empty list.
Line 4Once the source is drained, batch is empty and StopIteration is raised, so next(b, None) yields None.
Important notes
An exhausted iterator must keep raising StopIteration on every later call; code that revives it breaks the assumption for and list rely on.
__iter__ must return something that has __next__; returning a plain list or an object without __next__ makes the for statement raise TypeError, not a silent fallback.
Common mistakes
Ending __next__ with a bare return or return None instead of raise StopIteration, so the for loop never terminates and keeps handing you None forever.
Advancing the index before the bounds check, for example self.i += 1 then return self.data[self.i], which drops the first element and ends with IndexError instead of a clean stop.
Resetting the cursor inside __iter__ while still returning self, which makes a nested loop over the same object rewind the outer loop and spin forever.
Try it yourself
Change, predict, then run
Write an Evens class that takes a list and, via __iter__ and __next__, yields only the even numbers, then show that list() on the same instance returns [] the second time; refactor it into a container plus iterator pair so both calls give the same result.
Open the Python workspaceCheck your understanding
A class defines __iter__ that returns self and a stateful __next__. What happens when you nest two for loops over the same instance?
- The inner loop drains the remaining items, then the outer loop ends after its first item: one pass over the data in total
- Each loop gets its own cursor, so every outer item is paired with every inner item
- Python raises RuntimeError because the iterator is already being consumed
- The inner loop restarts from the beginning each time, because for always calls __iter__
Show answer
Both loops call iter() and get the same object back, so they share one cursor: the inner loop consumes everything left and the outer loop immediately sees StopIteration. Option 4 is tempting because for really does call __iter__, but __iter__ here returns self rather than a new cursor, so there is nothing to restart; only a container whose __iter__ builds a fresh iterator gives the behaviour in option 2.