Search code examples
pythonexceptiongenerator

Exception handling in Python generator function


I have to create Python function-generator, that should return all divisors of the positive number. If there are no divisors left function should return None. Done it this way:

def divisor(num):
    try:
        final = (x for x in range(1, num + 1) if num / x == int(num / x))
    except StopIteration:
        return None
    else:
        return final

Get such results:

three = divisor(3)
next(three) => 1
next(three) => 3

But when I call the function one more time, get StopIteration error:

next(three) =>
Traceback (most recent call last):
 ...
StopIteration

Why StopIteration is not handled? What am I doing wrong? How to fix it?


Solution

  • It happens because the function is not a generator. It is a regular function that returns the generator from the final variable on the first call. The best solution, in this case, is to transform the function into an actual generator and omit error handling at all:

    def divisor(num):
        for x in range(1, num + 1):
            if num / x == int(num / x):
                yield x
        while True:
            yield None