How to make an iterator in Python PREMIUM

How you can make your own iterator objects in Python?

Iterators need to work with iter and next

An iterable is an object that can be passed to the built-in iter function to get an iterator from it:

>>> numbers = [2, 1, 3, 4]
>>> iterator = iter(numbers)

And an iterator can then be passed to the built-in next function to get its next item:

>>> next(iterator)
2
>>> next(iterator)
1

But iterators are also iterables, which means you can pass them to the built-in iter function and they'll give you themselves back:

>>> iterator
<list_iterator object at 0x7fbb2a806ce0>
>>> iter(iterator)
<list_iterator object at 0x7fbb2a806ce0>

Let's try to create an iterator

Let's make a class whose objects are iterators.

We'll make a CountDown class which counts downward from a given number to zero. Here's the start of this class:

class CountDown:
    """Iterator which counts downward from a given number to 0."""

    def __init__(self, start):
        self.n = start

To make our class act as an iterator, we need to support the built-in next function and the built-in iter function.

The next function relies on __next__

The built-in next function relies on the __next__ method, so we need to implement a __next__ method.

Our __next__ method should count downward and also returning the given number:

    def __next__(self):
        number = self.n
        self.n -= 1
        return number

This __next__ method gets the current number, decrement its, and then returns the original number.

But right now, if we were to call next on this object, we would decrement downward forever. We need to raise a StopIteration exception when we reach 0:

class CountDown:
    """Iterator which counts downward from a given number to 0."""

    def __init__(self, start):
        self.n = start

    def __next__(self):
        if self.n <= 0:
            raise StopIteration("0 reached")
        number = self.n
        self.n -= 1
        return number

This CountDown class makes objects that work with the __next__ method:

>>> counter = CountDown(3)
>>> next(counter)
3
>>> next(counter)
2
>>> next(counter)
1

And this object hits 0, it'll raise a StopIteration exception:

>>> next(counter)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/trey/countdown.py", line 9, in __next__
StopIteration: 0 reached

But this class does not yet make fully functional iterators.

If we try to loop over this object, we would get an error:

>>> list(counter)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'CountDown' object is not iterable

All iterators should also be iterables. And our CountDown objects are not iterables yet. They're not iterables because they don't work with the built-in iter function.

The iter function relies on __iter__

To work with the built-in iter function, our objects need a __iter__ method. For iterators, the __iter__ method should return self, because __iter__ returns an iterator and iterators are already iterators:

class CountDown:
    """Iterator which counts downward from a given number to 0."""

    def __init__(self, start):
        self.n = start

    def __next__(self):
        if self.n <= 0:
            raise StopIteration("0 reached")
        number = self.n
        self.n -= 1
        return number

    def __iter__(self):
        return self

Now our CountDown class makes CountDown objects which work with a built-in next function:

>>> counter = CountDown(3)
>>> next(counter)
3

And which can also be looped over:

>>> list(counter)
[2, 1]

These objects are fully functional iterators now.

But... we probably shouldn't have made this class because there's an easier to make an iterator.

Avoid making iterator classes

Iterator classes are very rare to see.

The easiest way to make an iterator in Python is to make a generator object, either by using a generator expression or generator function.

This count_down generator function does basically the same thing that our CountDown iterator class did before:

def count_down(start):
    """Iterator which counts downward from a given number to 0."""
    n = start
    while n > 0:
        yield n
        n -= 1

We can take the generator objects that count_down returns:

>>> counter = count_down(3)
>>> counter
<generator object count_down at 0x7fcf71c64190>

And pass them to the built-in next function.

>>> next(counter)
3

Or we could loop over them:

>>> list(counter)
[2, 1]

This generator function does work a little bit differently, though.

Our CountDown class from before actually exposed an attribute called n:

>>> counter = CountDown(3)
>>> next(counter)
3
>>> counter.n
2

And we could change that attribute to change where our iterator continued counting from:

>>> counter.n = 5
>>> next(counter)
5

Whereas, our generator object doesn't expose any attributes. And if we try to assign an attribute to it, we would get an error:

>>> counter = count_down(3)
>>> next(counter)
3
>>> counter.n = 5
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'generator' object has no
 attribute 'n'

Examples of iterator classes

Iterator classes are pretty much only used if you need to make an iterator object that has other functionality beyond just iteration.

File objects are an example of an iterator with other functionality.

Files are iterators, so you can pass file objects to next, or you can loop over them:

>>> f = open("my_file.txt")
>>> next(f)
'This is line 1 in the file\n'

But you could also seek within files. Or close files:

>>> f.close()

Or check to see whether the file is closed:

>>> f.closed
True

Also, csv.reader returns iterator objects.

You can csv.reader objects to next:

>>> import csv
>>> reader = csv.reader(open("au-states.csv"))
>>> next(reader)
['State', 'Capital']

And of course you can loop over them.

But you can also check what line number you're currently on with the line_num attribute:

>>> reader.line_num
1

Summary

The most common way to make an iterator is to make a generator by making a generator expression or a generator function.

But if you need to make an iterator that does more than just iteration (that has additional functionality) you could make an iterator class with a __next__ method and a __iter__ method.