Iterators and Generators

Warm-up Exercise: Write a generator that will generate the powers of 2 from 1 up to 1024 (which is 2 to the 10th). Do it in three different ways: by making an iterator class, by making a generator with yield statements, and by using a generator expression.

class PowersOfTwoIterator():
    def __init__(self):
        self.power = -1

    def __iter__(self):
        return self

    def __next__(self):
        self.power += 1
        if self.power > 10:
            raise StopIteration
        else:
            return 2**self.power

def powers_of_two_gen():
    for power in range(11):
        yield 2**power

powers_of_two_exp = (2**power for power in range(11))

print([i for i in PowersOfTwoIterator()])
print([i for i in powers_of_two_gen()])
print([i for i in powers_of_two_exp])

[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]

Exercise 1: Write a generator all_n_multiples(n) that will be able to generate every positive multiple of n (in an infinite loop). Make sure to have a break statement when you test your code so that it isn't stuck running forever.

(solution below)

.

.

.

.

.

.

.

.

def all_n_multiples(n):
    num = 0
    while True:
        num += n
        yield num


for num in all_n_multiples(10):
    if num > 100:
        break
    print(num)

10
20
30
40
50
60
70
80
90
100

Exercise 2: Write a generator first that takes in a parameter n and a generator/iterator gen and yields the first n elements of gen (or all of it if it has less than n elements).

(solution below)

.

.

.

.

.

.

.

.

def first(n, gen):
    count = 0
    for elem in gen:
        count += 1
        if count > n:
            break
        yield elem

for num in first(5,all_n_multiples(10)):
    print(num)

10
20
30
40
50

Exercise 3: Modify the CheatingCoin class from last discussion to allow us to iterate through a CheatingCoin. This will return a sequence of flips until we've returned our cheating flip.

(solution below)

.

.

.

.

.

.

.

.

import random

class CheatingCoin:
    def __init__(self, n):
        self.__flips = 0
        self.__cheat_flip = n

    def flip(self):
        self.__flips += 1
        if self.__flips == self.__cheat_flip:
            self.__flips = 0
            return True
        else:
            return random.randint(0,1) == 1

    def __iter__(self):
        return self

    def __next__(self):
        # if we just call self.flip(), it's hard to know when we're supposed to stop.
        if self.__flips == self.__cheat_flip:
            self.__flips = 0
            raise StopIteration

        self.__flips += 1
        if self.__flips == self.__cheat_flip:
            return True
        else:
            return random.randint(0,1) == 1

coin = CheatingCoin(3)
for flip in coin:
    print(flip)

True
False
True

Exercise 4: Modify the above CheatingCoin class so that the iterator is a separate class, say CoinIterator.

(Hint: you'll need to pass in the __flips and __cheat_flip value in the iterator somehow, along with the coin itself.)

(solution below)

.

.

.

.

.

.

.

.

import random

class CheatingCoin:
    def __init__(self, n):
        self.__flips = 0
        self.__cheat_flip = n

    def flip(self):
        self.__flips += 1
        if self.__flips == self.__cheat_flip:
            self.__flips = 0
            return True
        else:
            return random.randint(0,1) == 1

    def __iter__(self):
        return CoinIterator(self, self.__flips, self.__cheat_flip)

class CoinIterator:
    def __init__(self, coin, flips, n):
        self.__coin = coin
        self.__flips = flips
        self.__n = n
    def __next__(self):
        # as before, if we just call self.flip(), it's hard to know when we're supposed to stop.
        if self.__flips == self.__n:
            raise StopIteration

        self.__flips += 1

        return self.__coin.flip()

coin = CheatingCoin(3)
for flip in coin:
    print(flip)

False
False
True