Search code examples
pythoniteratorgenerator

Best way to convert generator into iterator class


Consider the following dummy example:

def common_divisors_generator(n, m):

    # Init code
    factors_n = [i for i in range(1, n + 1) if n%i == 0]
    factors_m = [i for i in range(1, m + 1) if m%i == 0]

    # Iterative code
    for fn in factors_n:
        for fm in factors_m:
            if fn == fm:
                yield fn

# The next line is fast because no code is executed yet
cdg = common_divisors_generator(1537745, 373625435)
# Next line is slow because init code is executed on first iteration call
for g in cdg:
    print(g)

The init code, which takes a long time to compute, is executed once the generator has been iterated for the first time (as opposed to when the generator it is initialized). I would prefer that the init code it is executed as the generator is initialized.

For this purpose I convert the generator into an iterator class as follows:

class CommonDivisorsIterator(object):

    def __init__(self, n, m):
        # Init code
        self.factors_n = [i for i in range(1, n + 1) if n%i == 0]
        self.factors_m = [i for i in range(1, m + 1) if m%i == 0]

    def __iter__(self):
        return self

    def __next__(self):
        # Some Pythonic implementation of the iterative code above
        # ...
        return next_common_divisor

All ways I can think of implementing the __next__ method above are very cumbersome as compared to the simplicity of the iterative code in the generator with the yield keyword.

What would be the most Pythonic way of implementing the __next__ method in the iterator class?

Alternatively, how can I modify the the generator so that the init code is executed at init time?


Solution

  • In both cases (whether you use a function or a class), the solution is to split the implementation into two functions: a setup function and a generator function.

    Using yield in a function turns it into a generator function, which means that it returns a generator when it's called. But even without using yield, nothing's preventing you from creating a generator and returning it, like so:

    def common_divisors_generator(n, m):
        factors_n = [i for i in range(1, n + 1) if n%i == 0]
        factors_m = [i for i in range(1, m + 1) if m%i == 0]
    
        def gen():
            for fn in factors_n:
                for fm in factors_m:
                    if fn == fm:
                        yield fn
    
        return gen()
    

    And if you're using a class, there's no need to implement a __next__ method. You can just use yield in the __iter__ method:

    class CommonDivisorsIterator(object):
        def __init__(self, n, m):
            self.factors_n = [i for i in range(1, n + 1) if n%i == 0]
            self.factors_m = [i for i in range(1, m + 1) if m%i == 0]
    
        def __iter__(self):
            for fn in self.factors_n:
                for fm in self.factors_m:
                    if fn == fm:
                        yield fn