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.

`code = ''.join([prefix, str(num)])` is a really wordy way of writing `prefix + str(num)`, and honestly probably more readable as `f"{prefix}{num}` — juanpa.arrivillaga, Jun 19 '20 at 19:05
I assumed it'd be more efficient than using `+` and `f'{x}'`, no? — i_use_the_internet, Jun 19 '20 at 19:08
Why are you using `yield` in `generate_room_numbers`? Just use `return` and you don't need to use `next` — juanpa.arrivillaga, Jun 19 '20 at 19:08
No, it isn't, not for two strings. `''.join` is efficient for joining many strings in a container with a separator — juanpa.arrivillaga, Jun 19 '20 at 19:09
This appears to be a customer generator. Why can you not simply use the base generator expression as your loop iterator? Do you need to reset and reuse a very large generator sequence for your actual application? If not, a simple loop seems to be the better implementation. — Prune, Jun 19 '20 at 19:10
The whole thing can just be replaced by `[f'{flat} #{num}' for num in range(10) for flat in FLATS]` (or a generator expression if you prefer) — juanpa.arrivillaga, Jun 19 '20 at 19:16
@juanpa.arrivillaga The OP can correct me if I'm wrong, but I am assuming that `generate_nums` is meant to represent something which, for whatever reason, does indeed need to be a generator, and that the question is fundamentally about how to wrap a generator in a loop inside another generator, rather than whether the specific `generate_nums` in this application would be more efficiently implemented using `return`. — alani, Jun 19 '20 at 19:28
@alaniwi yes - the `generate_nums()` needs to be a generator because, in reality, it generates a really huge list — i_use_the_internet, Jun 19 '20 at 20:04

score 2 · Answer 1 · answered Jun 19 '20 at 19:08

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

alani · Accepted Answer · 2020-06-19T19:21:33.333

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)

Generator comprehension using another generator

2 Answers2