Iteration in Python

If an object, o, has no __iter__() method and you tell the interpreter to iterate over it, the interpreter initializes an internal variable to zero and repeatedly calls the object's __getitem__() method with successively higher values of the internal variable. From the point of view of the object, it's as though it were being manipulated by this code:

# Approximate equivalent of:
#    for val in o:
#        # [loop body]
intern = 0
while True:
    try:
        val = o[intern]
    except IndexError:
        break
    # [loop body]
    intern += 1

In fact, you can create your own classes whose instances can be iterated over in this way. All you need to do is provide a __getitem__(n) method that raises an IndexError exception when the value of n is too high. Suppose you wanted to implement fixed-length sequences of objects. You could define a function to create an appropriate sequence (list or tuple or string) with the required number of components in it (so fls("*", 12) would return "************", for example).

Alternatively, you could define an fls class, whose __init__() method had the same signature as the function above. Create fls.py as shown.

"""
Simple demonstration of the "old iteration protocol" - still available.
"""

class fls(object):

    def __init__(self, val, times):
        self.val = val
        self.count = times

    def __getitem__(self, n):
        if n >= self.count:
            raise IndexError("Object has no item %s %n")
        return self.val

thing = fls("*", 5)
for c in thing:
    print(c)

thing = fls(120,3)
for c in thing:
    print(c)

Save and run it. You see the following output:

*
*
*
*
*
120
120
120

So, for iteration purposes, you can see that the fls objects appear to act like other sequences, only with very boring behavior because all elements are constrained to be the same—the only value that __getitem__() ever returns is the one that was passed in to __init__(). But the main point is that you know a little more about Python's iteration mechanism. Now try a few other cases for yourself—use an interactive console session to create and test out some further fls objects interactively.

The iteration mechanism outlined above is all very well when you are iterating over numbered items in a sequence, but it does not naturally extend to collections like sets and dicts, which do not specify a natural ordering for their items. Dicts, in fact, do have a __getitem__() method, but it takes a key value and returns the appropriate item (assuming that a key with that value exists—if there is no such key, it raises a KeyError exception). Sets don't even have a __getitem__() method, since they are effectively "item-less dicts".

It was to overcome issues like this that the "new-style" iteration protocol was defined. You learned above that the interpreter will look for an __iter__() method on the objects that you iterate over. If it finds __iter__(), it uses it to create an iterator from the iterable you are iterating over.

The iterator will have a __next__() method—this is a requirement of the iteration protocol. Each time around the loop, the interpreter obtains the next value for the iterable by calling the iterator's __next__() method. Again, you can perhaps understand this more easily with an approximate Python equivalent to a for-loop over a new-style iterable:

# Approximate equivalent of:
#    for val in o:
#        # [loop body]
it = o.__iter__()
while True:
    try:
        val = it.__next__()
    except StopIteration:
        break
    # [loop body]

You may wonder why Python insists on creating a new object for each iteration: couldn't it just use the iterable directly? The answer to that question is "no": the iterator contains the current state of the iteration, and code that iterates over the same iterable twice is perfectly legal. Iterating over the same iterator, however, gives results that are not usually what you want. You can see this by playing with the interactive interpreter.

>>> lst = [1, 2, 3]
>>> dir(lst)
[..., '__getitem__', ..., '__init__', ...]
>>> for i in lst:
...     for j in lst:
...         print(i, j)
...
1 1
1 2
1 3
2 1
2 2
2 3
3 1
3 2
3 3
>>> li = lst.__iter__()
>>> dir(li)
[..., '__iter__', ..., '__next__', ...]
>>> for i in li:
...     for j in li:
...         print(i, j)
...
1 2
1 3
>>> lii = li.__iter__()
>>> li
<list_iterator object at 0x...>
>>> lii
<list_iterator object at 0x...>
>>> l2 = lst.__iter__()
>>> l2.__next__()
1
>>> l2.__next__()
2
>>> l2.__next__()
3
>>> l2.__next__()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

What's the difference between the first for loop and the second? Look at the dir() listing of lst, which is a list instance, and note that it has both __iter__() and __getitem__() methods. When the interpreter iterates over the list, calling its __iter__() method creates a new iterator, yielding a complete sequence of values, every time it encounters a for loop.

We then created a list iterator object by manually calling our list's __iter__() method. Note that the list iterator object also has an __iter__() method, and adds a __next__() method, but it lacks a __getitem__(). The __iter__() method of the iterator is rather different from that of the list, however:

def __iter__(self):
    return self

In other words, each time you iterate over a list (which is an iterable), the call to its __iter__() method creates a new iterator, which has its own independent state. The iterator's __iter__() method, however, does not create a new iterator, which means that the inner and outer loops are sharing the same iterator. This in turn means that by the time the outer loop is trying to begin its second iteration, the iterator has already been exhausted by the inner loop and (for the second time) raises the StopIteration exception.

The final few statements demonstrated this by manually going through the steps that the interpreter does when iterating over a list. We saw the l2 iterator produce three values on successive __next__() calls before raising a StopIteration exception. Normally, of course, the exception is caught internally by the logic of the for loop, and therefore does not become visible.

In summary, calling an iterable's __iter__() method creates an iterator that can be used to iterate over the iterable.

Now that you understand Python's iteration processes somewhat better, you may be wondering whether you can define your own iterable classes. The answer is "yes"! You will need to provide an __iter__() method (which can simply return self if you are implementing an iterator rather than a more general iterable: this is usually OK, since when you write an iterator class it is easy to create multiple instances, each having independent state). The __next__() method should return successive values until there are no more, at which point it should raise a StopIteration exception.

Rather than create an example now, we'll create it in the next section. First, we'll create a generator, and then we'll build an equivalent iterator.

The really convenient thing about generator functions is that they allow you to perform all sorts of complex calculations to produce the values in a sequence, but the code that consumes (makes use of) these values can be entirely separated from the generator that produces them. The values are consumed in a simple for loop—or any other similar iterative context in Python, such as a list comprehension.

Not only do they make your code simpler by separating out the production and consumption of sequences, but generators allow you to create sequence values one at a time, as they are consumed. There is no need to build a list or tuple to store them in, which means your programs will use less storage and operate more quickly (though these advantages do not really make much difference unless the number of objects becomes large).

Suppose you need to produce sequences determined by a list, but need to repeat the first list element once, the second twice, and so on. So given a list [2, 4, 6], the resulting sequence would be 2, 4, 4, 6, 6, 6. Let's write a generator that produces such sequences. First, though, we'll write tests to ensure that our generator function works. Create testgen.py as shown:

"""
testgen.py: simple test for a sequence generator
"""
import unittest
from gen123 import gen123

class TestGen(unittest.TestCase):

    def testEmpty(self):
        self.assertEqual(list(gen123([])), [], "Empty list does not give empty list")

    def test123(self):
        self.assertEqual(list(gen123([1])), [1], "[1] does not give [1]")
        self.assertEqual(list(gen123([1, 2])), [1, 2, 2])
        self.assertEqual(list(gen123([1, 2, 3])), [1, 2, 2, 3, 3, 3])

if __name__ == "__main__":
    unittest.main()

As usual, we start out with a simple stub function to make sure that the tests fail. Now, create gen123.py as shown:

"""
gen123.py: generate sequences from a base list, repeating
           each element one more time than the last
"""

def gen123(m):
    yield None

Save and run the test program:

FF
======================================================================
FAIL: test123 (__main__.TestGen)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "testgen.py", line 13, in test123
    self.assertEqual(list(gen123([1])), [1], "[1] does not give [1]")
AssertionError: [1] does not give [1]

======================================================================
FAIL: testEmpty (__main__.TestGen)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "testgen.py", line 10, in testEmpty
    self.assertEqual(list(gen123([])), [], "Empty list does not give empty list")
AssertionError: Empty list does not give empty list

----------------------------------------------------------------------
Ran 2 tests in 0.032s

FAILED (failures=2)

Now, let's see how it does with some real code in there.

"""
gen123.py: generate sequences from a base list, repeating
           each element one more time than the last
"""

def gen123(m):
    yield None
    
    n = 0
    for item in m:
        n += 1
        for i in range(n):
            yield item

Save and run the test program:

..
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

As you learned above, it is also possible to write classes that obey the iteration protocol. You will end this lesson by writing an iterator equivalent of the generator function above. Since you want it to perform exactly the same as the gen123 generators, you can use the same tests to verify its operation—that is one of the benefits of a test-driven environment! The new component should ideally be a "drop-in replacement" for the generator function. Create class123.py as shown:

"""
A simple iterator object specification.
"""

class gen123:

    def __init__(self, lst):
        "Initialize the iterator object."
        self.lst = lst
        self.itemno = 0
        self.count = 1

    def __iter__(self):
        "This object is not an iterable."
        return self

    def __next__(self):
        "Return the next value in the output sequence."
        if self.count > self.itemno:
            try:
                self.val = self.lst[self.itemno]
            except IndexError:
                raise StopIteration
            self.itemno += 1
            self.count = 1
        self.count += 1
        return self.val

This code is considerably more complex. This should not be surprising, because generator functions were devised to solve this type of problem cleanly and simply.

Instead of calling the generator function, the test routine will now call your iterator's class (which, you will notice, has the same name). This causes its __init__() method to be run, and the list of values is stored as an instance variable. Two other instance variables are initialized: one to keep track of which item is currently being output, and the other to keep track of how many times the current value has been produced.

All the magic, of course, takes place in the __next__() method. First it checks to see whether it is time to move to the next element of the value list (the item number and count are set up initially to ensure that this branch is actioned on the first call). If so, the val instance variable is retrieved.

If no more values are available, the method raises a StopIteration exception to terminate the loop. Note carefully that this action can be repeated—once the method starts to raise the exception, it should be raised for every subsequent call.

Once the correct item value is established, the count is incremented and the value is returned as the result of the call.

This code is about twice as long as that of the generator solution, and so you would probably choose to write a generator function for problems like this. But if you need close control over iterative behavior, you may end up needing to write your own iterators.

Testing the module is easy. Just make the following change to the test program:

"""
testgen.py: simple test for a list generator function
"""
import unittest
from class123 import gen123

class TestGen(unittest.TestCase):

    def testEmpty(self):
        self.assertEqual(list(gen123([])), [], "Empty list does not give empty list")

    def test123(self):
        self.assertEqual(list(gen123([1])), [1], "[1] does not give [1]")
        self.assertEqual(list(gen123([1, 2])), [1, 2, 2])
        self.assertEqual(list(gen123([1, 2, 3])), [1, 2, 2, 3, 3, 3])

if __name__ == "__main__":
    unittest.main()

Save and run the updated test program; you should see a successful result immediately, thereby giving strong evidence that the two implementations are equivalent.