Search⌘ K
AI Features
Home
Courses
Mastering the Art of Programming in Python 3

Iterators in Python

Explore how iterators work in Python by implementing __iter__ and __next__ magic methods. Understand how to create custom iterator classes, manage infinite loops, and use the StopIteration exception for controlled iteration. This lesson deepens your knowledge of Python's iteration protocol and compatibility differences between Python 2 and 3.

We'll cover the following...

In the previous lesson, we pondered over the decency of the for loop. Now let’s unveil it and see the magic behind it.

What is an iterator?

In Python, an iterator is an object that implements the __iter__ and __next__ dunder methods.

⚙️ Note: Dunder methods are methods having two prefix and suffix underscores in the method name. They are also known as magic methods. For example: __init__, __len__, __repr__, etc.

The __iter__() function returns an iterator object that goes through each element of the given object. Whereas, the __next__() function returns the next item in the sequence.

Objects that support the __iter__ and __next__ methods automatically work with loops.

Iterator from scratch

To create an iterator class, it must implement the __iter__() and __next() methods. Let’s create an iterator that returns numbers.

Python 3.10.4
class Counter:
  def __iter__(self):
    self.n = 1
    return self  # Returning an iterator object
 
  def __next__(self):
    x = self.n
    self.n += 1
    return x  # Returning next item in the sequence
count = Counter()
iterator = iter(count)
print(next(iterator))
print(next(iterator))
print(next(iterator))
print(next(iterator))

The above snippet makes it clear that the Counter will iterate over the sequence of infinite numbers starting from 1.

Stopping the iteration

The above executable will continue returning numbers forever if next() statements are used boundlessly, or if an infinite for loop is applied.

for x in iterator:
  print(x)

But, who wants to iterate forever? To halt the iteration, we can use the StopIteration exception by adding a termination condition in the __next__() method.

❌ Note: If we call the iterator after all the elements have been iterated, then StopIterationError is raised.

Now, let’s add a termination condition:

Python 3.10.4
class Counter:
  def __init__(self):
    self.n =1
  def __iter__(self):
    return self  # Returning an iterator object
 
  def __next__(self):
    if self.n <= 10:  # Termination condition
      x = self.n
      self.n += 1
      return x  # Returning next item in the sequence
    else:
      raise StopIteration # Stoping the iteration
count = Counter()
iterator = iter(count)
for x in iterator:
  print(x)

We add a termination condition at line 9. If that condition is False, it directly raises the exception StopIteration from line 14. So, now you see the loop itself breaks when n exceeds the value of 10.

🧐 You may have noticed something different. We can add the __init__() function and shift self.n (the initialization statement) from __iter__() to __init__(), leaving __iter__() with the return statement only.

Python 2.x compatibility

There’s a slight difference when implementing iterators in Python 2 and Python3.

In Python 3, the next value from the iterator is received via the __next__() method. But, in Python 2, it’s done with next() method (without any underscores).