Randomly sample between multiple generators?

user_7026806 0
Viewed 52

I'm trying to iterate over multiple generators randomly, and skip those that are exhausted by removing them from the list of available generators. However, the CombinedGenerator doesn't call itself like it should to switch generator. Instead it throws a StopIteration when the smaller iterator is exhausted. What am I missing?

The following works:

gen1 = (i for i in range(0, 5, 1))
gen2 = (i for i in range(100, 200, 1))

list_of_gen = [gen1, gen2]
print(list_of_gen)

list_of_gen.remove(gen1)
print(list_of_gen)

list_of_gen.remove(gen2)
print(list_of_gen)

where each generator is removed by their reference.

But here it doesn't:

import random

gen1 = (i for i in range(0, 5, 1))
gen2 = (i for i in range(100, 200, 1))

total = 105

class CombinedGenerator:
    def __init__(self, generators):
        self.generators = generators

    def __call__(self):
        generator = random.choice(self.generators)

        try:
            yield next(generator)
        except StopIteration:
            self.generators.remove(generator)
            if len(self.generators) != 0:
                self.__call__()
            else:
                raise StopIteration

c = CombinedGenerator([gen1, gen2])

for i in range(total):
    print(f"iter {i}")
    print(f"yielded {next(c())}")
python
2 Answers

As @Tomerikoo mentioned, you are basically creating your own Generator and it is better to implement __next__ which is cleaner and pythonic way.

The above code can be fixed with below lines.

def __call__(self):
    generator = random.choice(self.generators)

    try:
        yield next(generator)
    except StopIteration:
        self.generators.remove(generator)
        if len(self.generators) != 0:
            # yield your self.__call__() result as well
            yield next(self.__call__())
        else:
            raise StopIteration
user_66300120

First of all, in order to fix your current code, you just need to match the pattern you created by changing the line:

self.__call__()

to:

yield next(self.__call__())

Then, I would make a few small changes to your original code:

  • Instead of implementing __call__ and calling the object, it seems more reasonable to implement __next__ and simply call next on the object.
  • Instead of choosing the generator, I would choose the index. This mainly serves for avoiding the use of remove which is not so efficient when you can directly access the deleted object.
  • Personally I prefer to avoid recursion where possible so will change where I check that are still generators to use:
class CombinedGenerator:
    def __init__(self, generators):
        self.generators = generators

    def __next__(self):
        while self.generators:
            i = random.choice(range(len(self.generators)))

            try:
                return next(self.generators[i])
            except StopIteration:
                del self.generators[i]

        raise StopIteration

c = CombinedGenerator([gen1, gen2])

for i in range(total):
    print(f"iter {i}")
    print(f"yielded {next(c)}")

A nice bonus can be to add this to your class:

    def __iter__(self):
        return self

Which then allows you to directly iterate on the object itself and you don't need the total variable:

for i, num in enumerate(c):
    print(f"iter {i}")
    print(f"yielded {num}")
user_60458000
Related