CODEWITHSUNDEEP
pythonpython-3.xiteration
pstatix
pstatix

Reputation: 3848

Handling early __iter__ termination

I have a class that implements __iter__ and I have the following case:

for i in MyClass():
    if i == 'some condition':
        break
    # do other stuff

I need to be able to handle when iteration is terminated early. Can this be achieved within the classes implementation?

Upvotes: 1

Views: 66

Answers (1)

Holt
Holt

Reputation: 37626

You could try catching GeneratorExit:

class A:
    def __iter__(self):
        # You need to adapt this to your actual __iter__ implementation:
        i = 0
        try:
            while i < 5:
                yield i
                i += 1
        except GeneratorExit:
            print("GeneratorExit")

This will work if you do not store the generator, e.g. in

# If you break in any of these two loops, you'll get GeneratorExit printed:
a = A()
for i in a: ...
for i in A(): ...

If you store the generator, you'll get GeneratorExit only if the generator is garbage-collected without being exhausted:

a = A()
it = iter(a)
for i in it:
    if i == 2:
        break  # Does not print GeneratorExit.

it = iter(a)  # Prints GeneratorExit.

In any case, you probably don't want to clean-up if the generator is still alive, otherwise this would get weird:

it = iter(A())
for i in it:
    if i == 2:
        break

for i in it:  # What happens here if you cleaned-up?
    pass

Upvotes: 2

Related Questions