Question

I have a generator to be consumed by various consumers. Each of the latter can take different items from the generator, so I can't just use one big for-loop to take care of all the items. What I want is to completely consume the generator. How can it be done?

# -*- coding: utf-8 -*-
MEALS = ['Oysters', 'Consommé', 'Lamb', 'Rice', 'Sirloin','Banana', 'Pastry']

def server():
    for n in MEALS:
        yield n

def client(course, take):
    meal = []
    for _ in range(take):
        some_meal = next(course)
        meal.append(some_meal)
    return meal

if __name__ == '__main__':
    #print("Available meals: ", list(MEALS))
    course = server()
    try:
        while True:
            meal = client(course, 3)
            print("First client: ", meal)
            meal = client(course, 2)
            print("Second client: ", meal)
    except StopIteration:
        pass

Current output:

First client:  ['Oysters', 'Consommé', 'Lamb']
Second client:  ['Rice', 'Sirloin']

But where are the desserts??

Expected output:

First client:  ['Oysters', 'Consommé', 'Lamb']
Second client:  ['Rice', 'Sirloin']
First client:  ['Banana', 'Pastry']

UPDATE The accepted solution below with the added test on the returned list is OK except that I oversimplified the example code (There can be many next statements in client). What I now need is a way to return from the client function as soon as the first StopIteration is raised. So I added a follow-up question about the best way to exit a function upon hitting the first StopIteration.

Was it helpful?

Solution

In the second iteration of the while loop, the server generator only has 2 more items to yield, and the client() function will trigger the StopIteration exception when it tries to get 3 elements.

You'd need to handle the StopIteration in the client() function instead:

def client(course, take):
    meal = []
    for _ in range(take):
        try:
            some_meal = next(course)
            meal.append(some_meal)
        except StopIteration:
            pass
    return meal

Now that a client will handle the StopIteration, you'll have to handle the while loop differently; if a client() doesn't return elements your server must've been empty:

while True:
    meal = client(course, 3)
    if not meal:
        break
    print("First client: ", meal)
    meal = client(course, 2)
    print("Second client: ", meal)
    if not meal:
        break

You are missing a few tricks from the Python standard library here. You can reimplement your server with iter():

def server():
    return iter(MEALS)

and you could use itertools.islice() to handle your client:

from itertools import islice

def client(course, take):
    return list(islice(course, take))
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top