Peter Bengtsson: GeneratorExit - How to clean up after the last yield in Python

tl;dr; Use except GeneratorExit if your Python generator needs to know the consumer broke out.

Suppose you have a generator that yields things out. After each yield you want to execute some code that does something like logging or cleaning up. Here one such trivialized example:

The Problem

defpump():numbers=[1,2,3,4]fornumberinnumbers:yieldnumberprint("Have sent",number)print("Last number was sent")fornumberinpump():print("Got",number)print("All done")

The output is, as expected:

Got 1
Have sent 1
Got 2
Have sent 2
Got 3
Have sent 3
Got 4
Have sent 4
Last number was sent
All done

In this scenario, the consumer of the generator (the for number in pump() loop in this example) gets every thing the generator generates so after the last yield the generator is free to do any last minute activities which might be important such as closing a socket or updating a database.

Suppose the consumer is getting a bit "impatient" and breaks out as soon as it has what it needed.

defpump():numbers=[1,2,3,4]fornumberinnumbers:yieldnumberprint("Have sent",number)print("Last number was sent")fornumberinpump():print("Got",number)# THESE TWO NEW LINESifnumber>=2:breakprint("All done")

What do you think the out is now? I'll tell you:

Got 1
Have sent 1
Got 2
All done

In other words, the potentially important lines print("Have sent", number) and print("Last number was sent") never gets executed! The generator could tell the consumer (through documentation) of the generator "Don't break! If you don't want me any more raise a StopIteration". But that's not a feasible requirement.

The Solution

But! There is a better solution and that's to catch GeneratorExit exceptions.

defpump():numbers=[1,2,3,4]try:fornumberinnumbers:yieldnumberprint("Have sent",number)exceptGeneratorExit:print("Exception!")print("Last number was sent")fornumberinpump():print("Got",number)ifnumber==2:breakprint("All done")

Now you get what you might want:

Got 1
Have sent 1
Got 2
Exception!
Last number was sent
All done

Next Level Stuff

Note in the last example's output, it never prints Have sent 2 even though the generator really did send that number. Suppose that's an important piece of information, then you can reach that inside the except GeneratorExit block. Like this for example:

defpump():numbers=[1,2,3,4]try:fornumberinnumbers:yieldnumberprint("Have sent",number)exceptGeneratorExit:print("Have sent*",number)print("Last number was sent")fornumberinpump():print("Got",number)ifnumber==2:breakprint("All done")

And the output is:

Got 1
Have sent 1
Got 2
Have sent* 2
Last number was sent
All done

The * is just in case we wanted to distinguish between a break happening or not. Depends on your application.