Generator comprehension using another generator

Question

I am trying to write a method that returns a Generator . The end result of these two methods is to get a combination of two lists in the form: 'A #1', 'B #1', ..., 'F #9'

FLATS = ['A', 'B', 'C', 'D', 'E', 'F']

def generate_nums() -> Generator[str, None, None]:
    prefix = '#'
    for num in range(10):
        code = ''.join([prefix, str(num)])

        yield code

def generate_room_numbers() -> Generator[str, None, None]:
    room_nums = generate_nums()

    yield (' '.join([flat_name, room_num]) for room_num in room_nums for flat_name in FLATS)

if __name__ == "__main__":
    result = generate_room_numbers()
    result = next(result) # I hate this. How do I get rid of this?
    for room in result:
        print(room)

This gives me the correct outcome. Although, my annoyance is the line result = next(result) . Is there a better way to do this? I looked at this answer as well as the yield from syntax but I can barely understand generators enough as it is.

Answer 1

You could use a generator expression and a f-string:

FLATS = ['A', 'B', 'C', 'D', 'E', 'F']
room_numbers = (f'{letter} #{i}' for i in range(1, 10) for letter in FLATS)

for room in room_numbers:
    print(room)

Output:

A #1
B #1
C #1
.
.
.
D #9
E #9
F #9

Answer 2

It will be best if the yield statement is put inside an explicit loop rather than trying to yield a generator.

Your generate_room_numbers should look like this:

def generate_room_numbers():
    for flat_name in FLATS:
        room_nums = generate_nums()
        for room_num in room_nums:    
            yield (' '.join([flat_name, room_num]))

Note that the generate_nums() is called inside the flat_name loop, because you cannot repeatedly iterate over the same iterator that it returns; after iterating through it, it is exhausted and generate_nums will raise StopIteration every time (so that iterating produces an empty sequence).

(If generate_nums is expensive, then you could of course do nums = list(generate_nums()) outside the flat_name loop and then iterate over that inside the loop, but if this requires potentially a lot of memory, then it could defeat much of the point in using a generator in the first place.)

The rest of your code is unchanged except that the result = next(result) in the main code is removed, but for convenience, here is the whole thing:

FLATS = ['A', 'B', 'C', 'D', 'E', 'F']

def generate_nums():
    prefix = '#'
    for num in range(10):
        code = ''.join([prefix, str(num)])

        yield code

def generate_room_numbers():
    for flat_name in FLATS:
        room_nums = generate_nums()
        for room_num in room_nums:    
            yield (' '.join([flat_name, room_num]))

if __name__ == "__main__":
    result = generate_room_numbers()
    # result = next(result)  <<==== NOT NEEDED ANY MORE
    for room in result:
        print(room)