Python: Iterators, Generators and Comprehensions

Iterators

An iterator is an object that implements the next protocol and raises StopIteration when exhausted.

A list is an iterable object but not an iterator.

xs = [1,2,3]

try:
    next(xs)
except Exception as e:
    print(e)

'list' object is not an iterator

We can make an iterator out of an iterable object by calling the iter function.

xsi = iter(xs)

next(xsi)

1

next(xsi)

2

next(xsi)

3

try:
    next(xsi)
except Exception as e:
    print(type(e))

<class 'StopIteration'>

For loops

The for loop automatically constructs an iterator from an interator object, then repeatedly calls next until a StopIteration is raised.

for x in xs:
    print(x)

1
2
3

Looping over collections

Tuples and list preserve order

for x in ('a', 'b', 'c'):
    print(x)

a
b
c

for x in ['a', 'b', 'c']:
    print(x)

a
b
c

Sets do not preserve order

for x in set(['a', 'b', 'b', 'c']):
    print(x)

b
c
a

Dictionaries prserve order of entry (new)

d = {'c': 1, 'b': 2, 'a': 3}
for k in d:
    print(k)

c
b
a

for k in d.values():
    print(k)

1
2
3

for k, v in d.items():
    print(k, v)

c 1
b 2
a 3

Ranges

for i in range(3):
    print(i)

0
1
2

for i in range(1,4):
    print(i)

1
2
3

for i in range(1,6,2):
    print(i)

1
3
5

Enumerate

Traditional indexing

xs = ['a', 'b', 'c']

for i in range(len(xs)):
    print(i, xs[i])

0 a
1 b
2 c

Standard Python idiom is to use enumerate

for i, x in enumerate(xs):
    print(i, x)

0 a
1 b
2 c

for i, x in enumerate(xs, start=10):
    print(i, x)

10 a
11 b
12 c

Zip

list(zip(range(3), 'abc'))

[(0, 'a'), (1, 'b'), (2, 'c')]

for i, x in zip(range(3), 'abc'):
    print(i, x)

0 a
1 b
2 c

Generators

A generator is a function that produces a sequence of results instead of a single value. As generators do not store the sequence that is generated, they are memory-efficient. They use the yield and yield from keywords to return values.

Many built in Python functions return generators to minimize use of memory - e.g. range, zip, map, filter , open - you need to use a for loop to evaluate them.

zip(range(10), 'abcdefghij')

<zip at 0x104f778c8>

for i, c in zip(range(10), 'abcdefghij'):
    print(i, c)

0 a
1 b
2 c
3 d
4 e
5 f
6 g
7 h
8 i
9 j

Make sure that you have enough memory to do the conversion as a list stores all its contents in memory.

list(range(10))

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Generator functions

def updown(n):
    """Generator that goes up to n and down again."""
    for i in range(n):
        yield i
    for i in range(n, -1, -1):
        yield i

for i in updown(5):
    print(i, end=', ')

0, 1, 2, 3, 4, 5, 4, 3, 2, 1, 0,

def updown2(n):
    """Alternative version using yield from."""
    yield from range(n)
    yield from range(n, -1, -1)

for i in updown2(5):
    print(i, end=', ')

0, 1, 2, 3, 4, 5, 4, 3, 2, 1, 0,

def fib(n):
    """Return the next Fibonacci number up to n."""
    a, b = 1, 1
    for i in range(n):
        yield a
        a, b = b, a + b

for i in fib(10):
    print(i, end=', ')

1, 1, 2, 3, 5, 8, 13, 21, 34, 55,

Generator expressions

Syntax sugar for simple generators

def odd(n):
    """Traditional generator."""
    for i in range(n):
        if i %2 == 1:
            yield i

for i in odd(10):
    print(i, end=', ')

1, 3, 5, 7, 9,

Using a generator expression

gs = (i for i in range(10) if i%2 == 1)
for i in gs:
    print(i, end=', ')

1, 3, 5, 7, 9,

Comprehensions

Syntax sugar similar to that of generator expressions can be used to create lists, sets and dictionaries.

List comprehension

[i for i in range(10)]

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

[i**2 for i in range(10) if i % 2 == 1]

[1, 9, 25, 49, 81]

Set comprehension

{i for i in updown(5)}

{0, 1, 2, 3, 4, 5}

{i**3 for i in updown(5)}

{0, 1, 8, 27, 64, 125}

Dictionary comprehension

subjects = ['ann', 'bob', 'charles']
ages = [23, 34, 45]

{subject: age for subject, age in zip(subjects, ages)}

{'ann': 23, 'bob': 34, 'charles': 45}