4

I have a simple generator:

def counter(n): 
   counter = 0
   while counter <= n:
      counter += 1
      yield counter
   print(f"Nice! You counted to {counter}")

def test_counter(): 
   for count in counter(6): 
      print(f"count={count}")
      if count > 3: 
         break

In this example there is a break once count reaches 3. However the only problem with the code is that the break terminates the generator execution at the yield statement. It doesn't run the print statement after the yield statement. Is there a way to execute code within the generator that is after the yield statement when a break occurs?

Example run:

# What happens: 
>>> test_counter()
count=1
count=2
count=3
count=4

# What I'd like to see: 
>>> test_counter()
count=1
count=2
count=3
count=4
Nice! You counted to 4
LeanMan
  • 474
  • 1
  • 4
  • 18
  • Why not put that print statement inside `test_counter()`? Or, if that is indeed only a test function, in general just put the print statement after the call to `counter()`. You already have the `count` value anyways. – Camilo Martinez M. May 07 '21 at 06:14
  • I think you are missing the fact that this is a minimum example of the behavior I am looking for. I want to use it in a more complex function and for many more reasons than you stated. There is no test function and putting the print statement outside of the yield statement would only work if `counter` was accessible in `test_counter`. This is not possible since its a local variable of the generator (unless you know a way). Also `count` and `counter` are not the same values in memory. – LeanMan May 07 '21 at 06:22
  • Well, @Enzo and @Daniel have good solutions for you (no need for changing `print` statements location). But, just so it's out there, it is possible to access the value of count even after the break. Just try printing `count` after the for loop and you'll see what I mean. It will print 4 and that's what you would use to print your "Nice! You counted..." – Camilo Martinez M. May 07 '21 at 06:26
  • Yes that is accessing `count` but I need to be able to access the state of `counter`. Keep in mind that `yield` could return anything such as a tuple of values. This is just a simple example. Perhaps not a good one. – LeanMan May 07 '21 at 06:34
  • 1
    I understand, my bad. Some cool insights came out of your question, I am definitely using these answers in the future. – Camilo Martinez M. May 07 '21 at 06:36

2 Answers2

13

You have to put the print in a finally-block:

def counter(n): 
    try:
        counter = 0
        while counter <= n:
            counter += 1
            yield counter
    finally:
        print(f"Nice! You counted to {counter}")

def test_counter(): 
    for count in counter(6): 
        print(f"count={count}")
        if count > 3: 
            break
Daniel
  • 42,087
  • 4
  • 55
  • 81
  • 2
    This is really cool, but, a word of warning: the `finally` code only executes when the counter object is destroyed. I think in a `for` loop, that's when the loop ends (though I'm not 100% sure of that), but if you did `c = counter(6)` and then `for count in c`, you'd get the "Nice!" printed only when c was reassigned to something else, or maybe even never, if you weren't lucky. – joanis May 08 '21 at 21:41
  • PS: is it deliberate that your counter goes from 1 to n+1? E.g., `list(counter(3))` outputs `[1, 2, 3, 4]`. I guess that's a question for OP, since that part of the code is from them. – joanis May 08 '21 at 21:41
5

You can use send method of generator:

def counter(n): 
   counter = 0
   while counter <= n:
      counter += 1
      should_continue = yield counter
      if should_continue == False:
          break

   print(f"Nice! You counted to {counter}")

def test_counter(): 
    gen = counter(6)
    for count in gen: 
        print(f"count={count}")
        if count > 3: 
            try:
                gen.send(False)
            except StopIteration:
                break

Now test_counter() works as expected.

enzo
  • 9,861
  • 3
  • 15
  • 38