Writing/The Unbearable Lightness of Coding
Chapter 9 · 14 min read
Generators: from Iterators to Coroutines
I think we've talked about functions more than enough; so at long last, it's time to take another step up the Python ladder, which brings us to generators.
We'll start with a discussion on iteration, both recursive and lazy — first using plain functions, then higher-level ones, and finally generatos. We'll talk about the yield-statement, its interface and implementation, and see how it lets us save time and space, from infinite sequences to ephemeral comprehensions. We'll also address some issues with generators, like indexing, length, composition and exhaustion; clarify the difference between them and iterators; understand their prevalence; and gradually build towards full-fledged coroutines and the unusual design opportunities they present.
***
To start, let's develop some motivation. Suppose we'd like to iterate over a sequence — what do we do? If it's a simple list, we can use a for-loop; but what if it's something more dynamic, like the files and directories in our filesystem? Let's say we have the following directory structure:
dir/
file1
subdir/
file2
How do we scan the contents of dir — that is, dir/file1, dir/subdir, and dir/subdir/file2?
Recursive Iteration
If you're familiar with the standard library, you might know about os.listdir, which takes a directory and returns a list of its entries' names; and os.path.join, which concatenates paths correctly (using / on Linux and \ on Windows — doing this manually, with strings, wouldn't be cross-platform). Putting them together, we get:
>>> import os
>>> directory = 'dir'
>>> for name in os.listdir(directory):
... path = os.path.join(directory, entry)
... print(path)
dir/file1
dir/subdir
Unfortunately, this isn't enough; listdir doesn't recurse, so we don't go inside any subdirectories, and miss all the nested entries like dir/subdir/file2.
To traverse the entire directory tree, then, we'd have to use a while-loop:
>>> directories = ['dir']
>>> while directories:
... directory = directories.pop(0) # pop first item
... for name in os.listdir(directory):
... path = os.path.join(directory, name)
... print(path)
... if os.path.isdir(path):
... directories.append(path)
dir/file1
dir/subdir
dir/subdir/file2
Here, we maintain a list of all the directories we need to check; initially, it's just the top one, dir, where our search begins. As we go over its entries, we keep an eye out for subdirectories — so when we detect one, using os.path.isdir, we add it to the list of directories to consider.
This work well, so we can wrap it up in a function and use it in good health:
>>> def scan(directory):
... directories = [directory]
... paths = []
... while directories:
... directory = directories.pop(0)
... for name in os.listdir(directory):
... path = os.path.join(directory, name)
... paths.append(path)
... if os.path.isdir(path):
... directories.append(path)
... return paths
>>> for path in scan('dir'):
... print(path)
dir/file1
dir/subdir
dir/subdir/file2
As an aside, since we're using a function anyway, we can eliminate this ungainly list and use recursion instead: whenever we encounter a subdirectory, we can dive right in, combining the results of all the subcases:
>>> def scan(directory):
... paths = []
... for name in os.listdir(directory):
... path = os.path.join(directory, name)
... paths.append(path)
... if os.path.isdir(path):
... paths.extend(scan(path))
... return paths
>>> for path in scan('dir'):
... print(path)
dir/file1
dir/subdir
dir/subdir/file2
Both options are valid, and only differ in the order in which the entries are scanned. The first method uses breadth-first search, or BFS, returning all the paths on level before moving deeper into level , while the second one uses depth-first search, or DFS, delving as deep as possible whenever the opportunity presents itself. This is not immediately evident in our example, but a slightly more elaborate directory structure will illustrate the point:
dir/
subdir1/
file1
subdir2/
file2
Assuming that we rename our functions to scan_bfs and scan_dfs, we get:
>>> for path in scan_bfs('dir'):
... print(path)
dir/subdir1
dir/subdir2
dir/subdir1/file1
dir/subdir2/file2
>>> for path in scan_dfs('dir'):
... print(path)
dir/subdir1
dir/subdir1/file1
dir/subdir2
dir/subdir2/file2
Lazy Iteration
Like I said, both methods are fine; or at least, they work well, until we call either of them on a directory with lots of subdirectories. In this case, the way we've implemented our functions, they'd traverse the entire directory tree, loading all of its entries into memory, before the loop can even begin. This can be very slow and very wasteful, especially if we're just looking for a specific file, and would gladly abort the search as soon as we find it — so there's really no need to load everything in advance.
What else can we do? If we want a piecewise iteration, it means we have to keep some state: which pieces we've covered and which pieces are yet to come. Objects let us do this with ease, so we might be tempted to implement a DirectoryScan class; but I don't know, it feels like an overkill — as a rule of thumb, whenever we have a class with just one or two methods, we should ask ourselves if a simple function would suffice.
And besides, while it's true that iteration relies on a state, it's kind of incidental, isn't it? If we're counting from one to ten, the fact that we need to remember the previous number in order to determine the next one is an implementation detail, not some mutable configuration. The list of directories we kept to consider was also a state, but similarly incidental — it helped us orient ourselves, but was completely secondary to the essence of enumerating files and directories. This is evident from the fact that we were able to eliminate it when we switched to recursion — itself a form of incidental state, if you think about it: a call stack, implicitly maintained by the language.
So, if not a class, let's see if we can use a higher-order function instead; we've already seen how they're perfect to maintain precisely this sort of incidental state. For starters, let's do a simple iteration, counting from one integer to another. The state is obviously going to be a counter, and the trick would be to return not its value, but a function — which, when called, would provide us with the next value in the sequence. When it reaches its end, let's agree to return None:
>>> def integers(start, end):
... counter = start - 1
... def get_next():
... nonlocal counter
... if counter >= end:
... return None
... counter += 1
... return counter
... return get_next
Note that to modify the counter in the nested get_next function, we have to declare it nonlocal; and to guarantee that the first value is start, we have to start the counter from start - 1, so we can stick to a simple "increment-and-return" logic rather than treating the beginning as a special case. Anyway, here it is in action:
>>> get_next = integers(1, 3)
>>> get_next()
1
>>> get_next()
2
>>> get_next()
3
>>> get_next()
None
If you remember our chapter about loops, this None is what we called a sentinel value; so we can use the builtin iter function to fit this callable and its sentinel in a for-loop:
>>> for i in iter(integers(1, 3), None):
... print(i)
1
2
3
Maybe it's not the most elegant syntax imaginable, but consider what we've achieved: if we want to iterate over a million numbers, we don't have to create a list with all of them — we can use this function instead, which only keeps one number in memory at a time.
Lazy, Recursive Iteration
Now, we can implement scan in a similar fashion: load the contents of the top directory, but instead of recursing into all of its subdirectories immediately, return a function that pops them out one by one. This way, we only load additional entries that are one level deeper than where we're currently at, and only on demand; and when none remain, we return None like before:
>>> def scan(directory):
... paths = []
... for name in os.listdir(directory):
... paths.append(os.path.join(directory, name))
... def get_next():
... if not paths:
... return None
... path = paths.pop(0)
... if os.path.isdir(path):
... for name in os.listdir(path):
... paths.append(os.path.join(path, name))
... return path
... return get_next
This we can plug in a for-loop with a sentinel value:
>>> for path in iter(scan('dir'), None):
... print(path)
dir/subdir1
dir/subdir2
dir/subdir1/file1
dir/subdir2/file2
And rest easy that we don't strain the system or slow it down. Calling scan on a big directory tree is instantaneous, because it only loads the contents of the top directory and constructs the get_next function; and calling get_next is equally fast, since it loads the contents of one more subdirectory at most. What's more, we can stop calling it whenever we want, and nothing else gets loaded.
Generators
So we've successfully implemented a complex iterator using higher-level functions — but you have to admit, it's a bit of a pain: both on the definition side, where we need to work from a nested function, and on the iteration side, where we need to use iter with a sentinel.
Those difficulties led Python to a pretty fascinating syntax, which in turn led to an even more fascinating concept. It started with the idea of writing an iterator using a resumable function. You see, regular functions are invoked, run, and return a result — they happen sequentially, in a so called synchronous manner. But what if a function could pause in the middle, return some value — and later, resume from the same place? Maybe return another value, then pause again; maybe resume once more after that, and so on, until it's eventually finished? Such a function would be perfect complex iteration: it could use its local scope for exactly the kind of incidental state that's needed to keep track of it, and simply pause inbetween the values as it generates them.
To do this, we use yield instead of return. It's almost the same, expect the function doesn't end there and then — it just pauses, and can be resumed from this statement on at a later time. A function that has such a yield-statement is called a generator function:
>>> def gen():
... print('starting')
... yield 1
... print('resuming')
... yield 2
... print('ending')
And calling it doesn't actually execute anything yet — it just returns a generator object:
>>> gen()
<generator object gen at 0x…>
Iterating Generators
A generator object, then, is created in a paused state, positioned at the beginning of its flow and ready to go. To actually activate it, we need to pass it to the builtin next function, at which point we'll see our first message and get our first value:
>>> g = gen()
>>> next(g)
starting
1
Note that the generator is still very much alive — we can pass it to next again, and it will resume from the point where it had paused, giving us our second message and value:
>>> next(g)
resuming
2
If we call it again, we can expect to get our final message:
>>> next(g)
ending
StopIteration
But also a StopIteration, which a generator raises when it reaches its end.
At first glance, it's not a very elegant way to handle an iteration, either — but if you recall our chapter about loops once again, this is precisely Python's for protocol, so plugging a generator into a for-loop fits perfectly:
>>> for i in gen():
... print(i)
starting
1
resuming
2
ending
Implementing Generators
So on the iteration front, we're covered; let's see how the yield-statement plays out in implementation. Counting from one integer to another is as simple as:
>>> def integers(start, end):
... counter = start
... while counter <= end:
... yield counter
... counter += 1
>>> for i in integers(1, 3):
... print(i)
1
2
3
And note how we don't need to do any counter corrections — we can increment it after yielding, which wasn't possible with return. As for scanning a directory tree:
>>> def scan(directory):
... directories = [directory]
... while directories:
... directory = directories.pop(0)
... for name in os.listdir(directory):
... path = os.path.join(directory, name)
... yield path
... if os.path.isdir(path):
... directories.append(path)
>>> for path in scan('dir'):
... print(path)
dir/subdir1
dir/subdir1/file1
dir/subdir2
dir/subdir2/file2
Once again, it's more elegant than the original — this time, because we don't have to combine the results in a list, simply yield them as they come and move along.
Time and Space
The advantages that led us to consider lazy iteration — namely, saving execution time and memory space — naturally hold true for generators as well; and they're important enough and prevalent enough to repeat once again. Take this silly example, emulating a "heavy" iteration that might take 3 seconds to complete:
>>> import time
>>> def gen():
... time.sleep(1)
... yield 1
... time.sleep(1)
... yield 2
... time.sleep(1)
... yield 3
If we're only trying to check that something (e.g. 2) exists, we can bail out as soon as we have our answer (i.e. after 2 seconds). Not so if we were to treat it as a list, in which case the whole process would have to complete before we can follow up on it:
>>> if 2 in gen():
... print('found 2')
found 2 # takes 2 seconds
>>> if 2 in list(gen()):
... print('found 2')
found 2 # takes 3 seconds
To take a more realistic example — if we were looking for the passwd file on a Linux system, doing so with a generator would be relatively fast:
>>> for path in scan('/'):
... if path.endswith('/passwd'):
... print(path)
... break
/etc/passwd # fast
Simply because we can break as soon as we find it. Doing this with our previous implementation would take much longer, because it would have to traverse the entire filesystem before our loop would even begin.
Not only that, but loading the entire filesystem into a list might consume an enormous amount of memory, whereas a generator only needs to store what little information is necessary to compute the next result. Compare the size of a list with a million numbers:
>>> import sys
>>> x = list(range(1000000))
>>> sys.getsizeof(x)
8000056
Which is almost 8 megabytes, to that of our generator:
>>> sys.getsizeof(integers(1, 1000000))
104
Which manages to store the concept of counting up to a million in slightly more than 100 bytes.
Infinite Generators
In fact, we can use generators to "store", so to speak, a sequence that is literally infinite — definitely not something that will fit in any memory:
>>> def infinity():
... counter = 1
... while True:
... yield counter
... counter += 1
Which can be useful, for example, if we want to assign a unique, incremental ID to an unknown number of objects:
>>> counter = infinity()
>>> class A:
... def __init__(self):
... self.id = next(counter)
>>> a1 = A()
>>> a1.id
1
>>> a2 = A()
>>> a2.id
2
Alternatively, we can use it to generate a more interesting sequence, like an infinite stream of pseudorandom numbers. I like the Blum-Blum-Shub algorithm — mainly because it has such an excellent name:
>>> def blum_blum_shub(seed):
... modulo = 11 * 19 # or any other 2 prime numbers
... cursor = seed
... while True:
... cursor = cursor**2 % modulo
... yield cursor
Printing out the result is a bit tricky — because, well, it's an infinite sequence; and we can't slice a generator, as in print(blum_blum_shub[:10]). For now, I can demonstrate it with a clumsy loop, using enumerate to cap it at a particular index:
>>> for i, n in enumerate(blum_blum_shub(3)):
... print(n)
... if i >= 9:
... break
9
81
82
36
42
92
104
157
196
169
Generator Comprehensions
- Such a powerful concept we even have comprehensions
- sum, any/all
- Looking for a number
- Can't index it
- map and filter
Generator Indexing
- Next
- What if it's empty?
- Next(,None)
- What if it's not the first one?
- islice
Generator Length
- This is not surprising — it's the nature of generators
- sum
- ilen
Generator Compoision
- chain
- yield from
- recursive scan
Generator Exhaustion
- Might be tempting to cast to list instead of remembering all of this, but it's dangerous
- Sometimes useful, like zipping a generator with itself to get pairs
Iterables, Iterators, Generators
- This is all reminiscent of iterators
- Clearing up terminology
Iterator Prevalence
- This is important, because iterators/generators are everywhere
- os.walk, itertools.count, range
- This wasn't always the case: xrange
- reversed, enumerate, zip — implementation with generators
- Why aren't they generators, then? — some historical, some deliberate object attributes/methods
Coroutines
- But generators are not just iterators
- Syntax and examples of send and throw
- Asynchronous programming