Rao Discussion 6: Iterators, Generators
This discussion worksheet is for the Rao offering of CS 61A. Your work is not graded and you do not need to submit anything.
Iterators
An iterable is an object whose elements we can go through one at a time.
Iterables can be used in for loops and list comprehensions. Iterables can
also be converted into lists using the list
function.
Examples of iterables we have seen so far in Python include strings,
lists, tuples, and ranges.
>>> for x in "cat":
... print(x)
c
a
t
>>> [x*2 for x in (1, 2, 3)]
[2, 4, 6]
>>> list(range(4))
[0, 1, 2, 3]
Note the abstraction present here: given some iterable object, we have a set of defined actions that we can take with it. In this discussion, we will peak below the abstraction barrier and examine how iterables are implemented "under the hood".
In Python, iterables are formally implemented as objects that can be passed into the built-in iter
function to produce an iterator. An iterator is another type of object
that can produce elements one at a time with the next
function.
iter(iterable)
: Returns an iterator over the elements of the given iterable.next(iterator)
: Returns the next element in an iterator, or raises aStopIteration
exception if there are no elements left.
For example, a list of numbers is an iterable,
since iter
gives us an iterator over the given sequence, which
we can navigate using the next
function:
>>> lst = [1, 2, 3]
>>> lst_iter = iter(lst)
>>> lst_iter
<list_iterator object ...>
>>> next(lst)
1
>>> next(lst)
2
>>> next(lst)
3
>>> next(lst)
StopIteration
Iterators are very simple. There is only a mechanism to get the next element
in the iterator: next
. There is no way to index into an iterator
and there is no way to go backward. Once an iterator has produced an element,
there is no way for us to get that element again unless we store it.
Note that iterators themselves are iterables: calling iter
on an iterator
simply returns the same iterator object.
For example, we can see what happens when we use iter
and next
with a list:
>>> lst = [1, 2, 3]
>>> next(lst) # Calling next on an iterable
TypeError: 'list' object is not an iterator
>>> list_iter = iter(lst) # Creates an iterator for the list
>>> next(list_iter) # Calling next on an iterator
1
>>> next(iter(list_iter)) # Calling iter on an iterator returns itself
2
>>> for e in list_iter: # Exhausts remainder of list_iter
... print(e)
3
>>> next(list_iter) # No elements left!
StopIteration
>>> lst # Original iterable is unaffected
[1, 2, 3]
The map
and filter
functions we learned earlier in class return iterator objects.
Q1: WWPD: Iterators
What would Python display?
>>> s = "cs61a"
>>> s_iter = iter(s)
>>> next(s_iter)
>>> next(s_iter)
>>> list(s_iter)
>>> s = [[1, 2, 3, 4]]
>>> i = iter(s)
>>> j = iter(next(i))
>>> next(j)
>>> s.append(5)
>>> next(i)
>>> next(j)
>>> list(j)
>>> next(i)
Generators
We can define custom iterators by writing a generator function, which returns a special type of iterator called a generator.
A generator function looks like a normal Python function, except
that it has at least one yield
statement.
When we call a generator function, a generator object
is returned without evaluating the body of the generator function itself.
(Note that this is different from ordinary Python functions.
While generator functions and normal functions look the same,
their evaluation rules are very different!)
When we first call next
on the returned generator,
we will begin evaluating the body of the generator function until
an element is yielded or the function otherwise stops
(such as if we return
).
The generator remembers where we stopped,
and will continue evaluating from that stopping point
on the next time we call next
.
As with other iterators, if there are no more elements to be generated,
then calling next
on the generator will give us a StopIteration
.
For example, here's a generator function:
def countdown(n):
print("Beginning countdown!")
while n >= 0:
yield n
n -= 1
print("Blastoff!")
To create a new generator object, we can call the generator function.
Each returned generator object from a function call will separately
keep track of where it is in terms of evaluating the body of the function.
Like all other iterators, calling iter
on an existing generator object
returns the same generator object.
>>> c1, c2 = countdown(2), countdown(2)
>>> c1 is iter(c1) # a generator is an iterator
True
>>> c1 is c2
False
>>> next(c1)
Beginning countdown!
2
>>> next(c2)
Beginning countdown!
2
In a generator function, we can also have a yield from
statement,
which will yield each element from an iterator or iterable.
>>> def gen_list(lst):
... yield from lst
...
>>> g = gen_list([1, 2])
>>> next(g)
1
>>> next(g)
2
>>> next(g)
StopIteration
Since generators are themselves iterators, this means we can use yield from
to create recursive generators!
>>> def rec_countdown(n):
... if n < 0:
... print("Blastoff!")
... else:
... yield n
... yield from rec_countdown(n-1)
...
>>> r = rec_countdown(2)
>>> next(r)
2
>>> next(r)
1
>>> next(r)
0
>>> next(r)
Blastoff!
StopIteration
Q2: WWPD: Generators
What would Python display? If the command errors, input the specific error.
>>> def infinite_generator(n):
... yield n
... while True:
... n += 1
... yield n
>>> next(infinite_generator)
>>> gen_obj = infinite_generator(1)
>>> next(gen_obj)
>>> next(gen_obj)
>>> list(gen_obj)
>>> def rev_str(s):
... for i in range(len(s)):
... yield from s[i::-1]
>>> hey = rev_str("hey")
>>> next(hey)
>>> next(hey)
>>> next(hey)
>>> list(hey)
>>> def add_prefix(s, pre):
... if not pre:
... return
... yield pre[0] + s
... yield from add_prefix(s, pre[1:])
>>> school = add_prefix("schooler", ["pre", "middle", "high"])
>>> next(school)
>>> list(map(lambda x: x[:-2], school))
Q3: Filter-Iter
Implement a generator function called filter_iter(iterable, f)
that only yields
elements of iterable
for which f
returns True
.
Run in 61A CodeRemember,
iterable
could be infinite!
Q4: What's the Difference?
Implement differences
, a generator function that takes an iterable it
whose elements are numbers.
differences
should produce a generator that yield the differences between successive terms of it
.
If it
has less than 2 values, differences
should yield nothing.
Q5: Primes Generator
Write a function primes_gen
that takes a single argument n
and yields all prime
numbers less than or equal to n
in decreasing order. Assume n >= 1
.
You may use the is_prime
function included below, which we implemented in
Discussion 3.
First approach this problem using a for
loop and using yield
.
Now that you've done it using a for
loop and yield
, try using yield from
!
Optional Challenge: Now rewrite the generator so that it also prints the primes in ascending order.
Run in 61A CodeQ6: Stair Ways
In discussion 4, we considered how many different ways there are to climb a flight of stairs with n
steps if you are able to take 1 or 2 steps at a time. In this problem, you will write a generator function stair_ways
that yields all the different ways you can climb such a staircase.
Each "way" of climbing a staircase is represented by a list of 1
s and 2
s, representing the sequence of step sizes a person should take to climb the flight.
For example, for a flight with 3 steps, there are three ways to climb it:
- You can take one step each time:
[1, 1, 1]
. - You can take two steps then one step:
[2, 1]
. - You can take one step then two steps:
[1, 2].
.
Therefore, stair_ways(3)
should yield [1, 1, 1]
, [2, 1]
, and [1, 2]
in any order.