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
def pump(): numbers = [1, 2, 3, 4] for number in numbers: yield number print("Have sent", number) print("Last number was sent") for number in pump(): 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.
def pump(): numbers = [1, 2, 3, 4] for number in numbers: yield number print("Have sent", number) print("Last number was sent") for number in pump(): print("Got", number) # THESE TWO NEW LINES if number >= 2: break print("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.
def pump(): numbers = [1, 2, 3, 4] try: for number in numbers: yield number print("Have sent", number) except GeneratorExit: print("Exception!") print("Last number was sent") for number in pump(): print("Got", number) if number == 2: break print("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:
def pump(): numbers = [1, 2, 3, 4] try: for number in numbers: yield number print("Have sent", number) except GeneratorExit: print("Have sent*", number) print("Last number was sent") for number in pump(): print("Got", number) if number == 2: break print("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.
Comments
Related posts
- Previous:
- Writing a custom Datadog reporter using the Python API 21 May 2018
- Related by Keyword:
- Musings about django.contrib.auth.models.User 28 August 2010
- Related by Text:
- jQuery and Highslide JS 08 January 2008
- I'm back! Peterbe.com has been renewed 05 June 2005
- Anti-McCain propaganda videos 12 August 2008
- I'm Prolog 01 May 2007
- Ever wondered how much $87 Billion is? 04 November 2003