Lab 6: Iterators, Generators

Due by 11:59pm on Monday, July 15.

Starter Files

Download lab06.zip. Inside the archive, you will find starter files for the questions in this lab, along with a copy of the Ok autograder.

Topics

Here's a refresher on Iterators and Generators. It's okay to skip directly to the questions and refer back here if you get stuck.

Iterators

An iterable is any value that can be iterated through, or gone through one element at a time. One construct that we've used to iterate through an iterable is a for statement:
for elem in iterable:
    # do something

In general, an iterable is an object on which calling the built-in iter function returns an iterator. An iterator is an object on which calling the built-in next function returns the next value.

For example, a list is an iterable value.

>>> s = [1, 2, 3, 4]
>>> next(s)       # s is iterable, but not an iterator
TypeError: 'list' object is not an iterator
>>> t = iter(s)   # Creates an iterator
>>> t
<list_iterator object ...>
>>> next(t)       # Calling next on an iterator
1
>>> next(t)       # Calling next on the same iterator
2
>>> next(iter(t)) # Calling iter on an iterator returns itself
3
>>> t2 = iter(s)
>>> next(t2)      # Second iterator starts at the beginning of s
1
>>> next(t)       # First iterator is unaffected by second iterator
4
>>> next(t)       # No elements left!
StopIteration
>>> s             # Original iterable is unaffected
[1, 2, 3, 4]

You can also use an iterator in a for statement because all iterators are iterable. But note that since iterators keep their state, they're only good to iterate through an iterable once:

>>> t = iter([4, 3, 2, 1])
>>> for e in t:
...     print(e)
4
3
2
1
>>> for e in t:
...     print(e)

There are built-in functions that return iterators.

>>> m = map(lambda x: x * x, [3, 4, 5])
>>> next(m)
9
>>> next(m)
16
>>> f = filter(lambda x: x > 3, [3, 4, 5])
>>> next(f)
4
>>> next(f)
5
>>> z = zip([30, 40, 50], [3, 4, 5])
>>> next(z)
(30, 3)
>>> next(z)
(40, 4)

Generators

We can create our own custom iterators by writing a generator function, which returns a special type of iterator called a generator. Generator functions have yield statements within the body of the function instead of return statements. Calling a generator function will return a generator object and will not execute the body of the function.

For example, let's consider the following generator function:

def countdown(n):
    print("Beginning countdown!")
    while n >= 0:
        yield n
        n -= 1
    print("Blastoff!")

Calling countdown(k) will return a generator object that counts down from k to 0. Since generators are iterators, we can call iter on the resulting object, which will simply return the same object. Note that the body is not executed at this point; nothing is printed and no numbers are outputted.

>>> c = countdown(5)
>>> c
<generator object countdown ...>
>>> c is iter(c)
True

So how is the counting done? Again, since generators are iterators, we call next on them to get the next element! The first time next is called, execution begins at the first line of the function body and continues until the yield statement is reached. The result of evaluating the expression in the yield statement is returned. The following interactive session continues from the one above.

>>> next(c)
Beginning countdown!
5

Unlike functions we've seen before in this course, generator functions can remember their state. On any consecutive calls to next, execution picks up from the line after the yield statement that was previously executed. Like the first call to next, execution will continue until the next yield statement is reached. Note that because of this, Beginning countdown! doesn't get printed again.

>>> next(c)
4
>>> next(c)
3

The next 3 calls to next will continue to yield consecutive descending integers until 0. On the following call, a StopIteration error will be raised because there are no more values to yield (i.e. the end of the function body was reached before hitting a yield statement).

>>> next(c)
2
>>> next(c)
1
>>> next(c)
0
>>> next(c)
Blastoff!
StopIteration

Separate calls to countdown will create distinct generator objects with their own state. Usually, generators shouldn't restart. If you'd like to reset the sequence, create another generator object by calling the generator function again.

>>> c1, c2 = countdown(5), countdown(5)
>>> c1 is c2
False
>>> next(c1)
5
>>> next(c2)
5

Here is a summary of the above:

  • A generator function has a yield statement and returns a generator object.
  • Calling the iter function on a generator object returns the same object without modifying its current state.
  • The body of a generator function is not evaluated until next is called on a resulting generator object. Calling the next function on a generator object computes and returns the next object in its sequence. If the sequence is exhausted, StopIteration is raised.
  • A generator "remembers" its state for the next next call. Therefore,

    • the first next call works like this:

      1. Enter the function and run until the line with yield.
      2. Return the value in the yield statement, but remember the state of the function for future next calls.
    • And subsequent next calls work like this:

      1. Re-enter the function, start at the line after the yield statement that was previously executed, and run until the next yield statement.
      2. Return the value in the yield statement, but remember the state of the function for future next calls.
  • Calling a generator function returns a brand new generator object (like calling iter on an iterable object).
  • A generator should not restart unless it's defined that way. To start over from the first element in a generator, just call the generator function again to create a new generator.

Another useful tool for generators is the yield from statement. yield from will yield all values from an iterator or iterable.

>>> def gen_list(lst):
...     yield from lst
...
>>> g = gen_list([1, 2, 3, 4])
>>> next(g)
1
>>> next(g)
2
>>> next(g)
3
>>> next(g)
4
>>> next(g)
StopIteration

Getting Started Videos

These videos may provide some helpful direction for tackling the coding problems on this assignment.

To see these videos, you should be logged into your berkeley.edu email.

YouTube link

Required Questions

Q1: WWPD: Iterators

Important: Enter StopIteration if a StopIteration exception occurs, Error if you believe a different error occurs, and Iterator if the output is an iterator object.

Important: Python's built-in function map, filter, and zip return iterators, not lists.

Use Ok to test your knowledge with the following "What Would Python Display?" questions:

python3 ok -q iterators-wwpd -u

>>> s = [1, 2, 3, 4]
>>> t = iter(s)
>>> next(s)
______
Error
>>> next(t)
______
1
>>> next(t)
______
2
>>> next(iter(s))
______
1
>>> next(iter(s))
______
1
>>> u = t >>> next(u)
______
3
>>> next(t)
______
4
>>> r = range(6)
>>> r_iter = iter(r)
>>> next(r_iter)
______
0
>>> [x + 1 for x in r]
______
[1, 2, 3, 4, 5, 6]
>>> [x + 1 for x in r_iter]
______
[2, 3, 4, 5, 6]
>>> next(r_iter)
______
StopIteration
>>> map_iter = map(lambda x : x + 10, range(5))
>>> next(map_iter)
______
10
>>> next(map_iter)
______
11
>>> list(map_iter)
______
[12, 13, 14]
>>> for e in filter(lambda x : x % 4 == 0, range(1000, 1008)): ... print(e)
______
1000 1004
>>> [x + y for x, y in zip([1, 2, 3], [4, 5, 6])]
______
[5, 7, 9]

Q2: Count Occurrences

Implement count_occurrences, which takes an iterator t, an integer n, and a value x. It returns the number of elements in the first n elements of t that equal to x.

Important: Call next on t exactly n times. Assume there are at least n elements in t.

def count_occurrences(t, n, x):
    """Return the number of times that x is equal to one of the
    first n elements of iterator t.

    >>> s = iter([10, 9, 10, 9, 9, 10, 8, 8, 8, 7])
    >>> count_occurrences(s, 10, 9)
    3
    >>> t = iter([10, 9, 10, 9, 9, 10, 8, 8, 8, 7])
    >>> count_occurrences(t, 3, 10)
    2
    >>> u = iter([3, 2, 2, 2, 1, 2, 1, 4, 4, 5, 5, 5])
    >>> count_occurrences(u, 1, 3)  # Only iterate over 3
    1
    >>> count_occurrences(u, 3, 2)  # Only iterate over 2, 2, 2
    3
    >>> list(u)                     # Ensure that the iterator has advanced the right amount
    [1, 2, 1, 4, 4, 5, 5, 5]
    >>> v = iter([4, 1, 6, 6, 7, 7, 6, 6, 2, 2, 2, 5])
    >>> count_occurrences(v, 6, 6)
    2
    """
    "*** YOUR CODE HERE ***"

Use Ok to test your code:

python3 ok -q count_occurrences

Q3: Hailstone

Write a generator function that outputs the hailstone sequence starting at number n.

Here's a quick reminder of how the hailstone sequence is defined:

  1. Pick a positive integer n as the start.
  2. If n is even, divide it by 2.
  3. If n is odd, multiply it by 3 and add 1.
  4. Continue this process until n is 1.

Make sure you don't create an infinite generator!

As an extra challenge, try writing a solution using recursion. Since hailstone returns a generator, you can yield from a call to hailstone!

def hailstone(n):
    """Yields the elements of the hailstone sequence starting at n.

    >>> for num in hailstone(10):
    ...     print(num)
    10
    5
    16
    8
    4
    2
    1
    """
    "*** YOUR CODE HERE ***"

Use Ok to test your code:

python3 ok -q hailstone

Check Your Score Locally

You can locally check your score on each question of this assignment by running

python3 ok --score

This does NOT submit the assignment! When you are satisfied with your score, submit the assignment to Gradescope to receive credit for it.

Submit Assignment

Submit this assignment by uploading any files you've edited to the appropriate Gradescope assignment. Lab 00 has detailed instructions.

In addition, all students who are not in the mega lab must submit the attendance form for credit. Ask your section TA for the link! Submit this form for each section, whether you attended lab or missed it for a good reason. The attendance form is not required for mega section students.

Optional Questions

Q4: Merge

Implement merge(incr_a, incr_b), which takes two iterables incr_a and incr_b whose elements are ordered. merge yields elements from incr_a and incr_b in sorted order, eliminating repetition. You may assume incr_a and incr_b themselves do not contain repeats, and that none of the elements of either are None. You may not assume that the iterables are finite; either may produce an infinite stream of results.

You will probably find it helpful to use the two-argument version of the built-in next function: next(incr, v) is the same as next(incr), except that instead of raising StopIteration when incr runs out of elements, it returns v.

See the doctest for examples of behavior.

def merge(incr_a, incr_b):
    """Yield the elements of strictly increasing iterables incr_a and incr_b, removing
    repeats. Assume that incr_a and incr_b have no repeats. incr_a or incr_b may or may not
    be infinite sequences.

    >>> m = merge([0, 2, 4, 6, 8, 10, 12, 14], [0, 3, 6, 9, 12, 15])
    >>> type(m)
    <class 'generator'>
    >>> list(m)
    [0, 2, 3, 4, 6, 8, 9, 10, 12, 14, 15]
    >>> def big(n):
    ...    k = 0
    ...    while True: yield k; k += n
    >>> m = merge(big(2), big(3))
    >>> [next(m) for _ in range(11)]
    [0, 2, 3, 4, 6, 8, 9, 10, 12, 14, 15]
    """
    iter_a, iter_b = iter(incr_a), iter(incr_b)
    next_a, next_b = next(iter_a, None), next(iter_b, None)
    "*** YOUR CODE HERE ***"

Use Ok to test your code:

python3 ok -q merge

Q5: Deep Map

Write a function deep_map that takes a list s and a one-argument function f. s may be a nested list, one that contain other lists. deep_map modifies s by replacing each element within s or any of the lists it contains with the result of calling f on that element.

deep_map returns None and should not create any new lists.

Hint: type(a) == list will evaluate to True if a is a list.

def deep_map(f, s):
    """Replace all non-list elements x with f(x) in the nested list s.

    >>> six = [1, 2, [3, [4], 5], 6]
    >>> deep_map(lambda x: x * x, six)
    >>> six
    [1, 4, [9, [16], 25], 36]
    >>> # Check that you're not making new lists
    >>> s = [3, [1, [4, [1]]]]
    >>> s1 = s[1]
    >>> s2 = s1[1]
    >>> s3 = s2[1]
    >>> deep_map(lambda x: x + 1, s)
    >>> s
    [4, [2, [5, [2]]]]
    >>> s1 is s[1]
    True
    >>> s2 is s1[1]
    True
    >>> s3 is s2[1]
    True
    """
    "*** YOUR CODE HERE ***"

Use Ok to test your code:

python3 ok -q deep_map

Q6: Buying Fruit

Implement buy, which takes a list of required_fruits (strings), a dictionary prices (strings for key, positive integers for value), and a total_amount (integer). It prints all the ways to buy some of each required fruit so that the total price equals total_amount. You must include at least one of every fruit in required_fruit and cannot include any other fruits that are not in required_fruit.

The display function will be helpful. You can call display on a fruit and its count to create a string containing both.

What does fruits and amount represent? How are they used in the recursive?

def buy(required_fruits, prices, total_amount):
    """Print ways to buy some of each fruit so that the sum of prices is amount.

    >>> prices = {'oranges': 4, 'apples': 3, 'bananas': 2, 'kiwis': 9}
    >>> buy(['apples', 'oranges', 'bananas'], prices, 12)
    [2 apples][1 orange][1 banana]
    >>> buy(['apples', 'oranges', 'bananas'], prices, 16)
    [2 apples][1 orange][3 bananas]
    [2 apples][2 oranges][1 banana]
    >>> buy(['apples', 'kiwis'], prices, 36)
    [3 apples][3 kiwis]
    [6 apples][2 kiwis]
    [9 apples][1 kiwi]
    """
    def add(fruits, amount, cart):
        if fruits == [] and amount == 0:
            print(cart)
        elif fruits and amount > 0:
            fruit = fruits[0]
            price = ____
            for k in ____:
                add(____, ____, ____)
    add(required_fruits, total_amount, '')

def display(fruit, count):
    """Display a count of a fruit in square brackets.

    >>> display('apples', 3)
    '[3 apples]'
    >>> display('apples', 1)
    '[1 apple]'
    """
    assert count >= 1 and fruit[-1] == 's'
    if count == 1:
        fruit = fruit[:-1]  # get rid of the plural s
    return '[' + str(count) + ' ' + fruit + ']'

Use Ok to test your code:

python3 ok -q buy