Clean up on for loop break

Admin 1
Viewed 591

Can you create an iterable in python which runs clean up code when for loop exits? Something like:

from random import randint

class Iterable:
    def __iter__(self):
        return self
    def __next__(self):
        return randint(1, 10)
    def __iterclose__(self):
        print("Clean up code")

for x in Iterable():
    if x < 5:
        break

# Prints "Clean up code"
iterable iterator for-loop python-3.x python
2 Answers

I think in order to benefit from the iterable protocol, you should let your Iterable gain control about the stop condition. How about this:


from random import randint


class Iterable:

    def __init__(self, max):
        self.i = 0
        self.max = max

    def __iter__(self):
        return self

    def __next__(self):
        self.i += 1
        if self.i > self.max:
            self.cleanup()
            raise StopIteration
        return randint(1, 10)

    def cleanup(self):
        print("Clean up code")


for x in Iterable(5):
    print(x)
user_76634100

As I said in the comment :

When finishing iteration, you should raise the StopIteration exception, at this point you should already know that you have to perform the clean up, else this iterator will run forever.

Also iterator and iterable are different things, from fluent python book:

iterator

Any object that implements the next no-argument method that returns the next item in a series or raises StopIteration when there are no more items. Python iterators also implement the iter method so they are iterable as well

Where:

iterable Any object from which the iter built-in function can obtain an iterator. Objects implementing an iter method returning an iterator are iterable.

So iterator should know when it will end, Hopefully this code will clear things up:

from random import randint

class MyIterator:
  def __init__(self, n):
    self.n = n
    self.i = 0

  def __iter__(self):
    return self

  def __next__(self):
    if self.i < self.n:
      self.i += 1
      return randint(1, 10)
    else:
      print("doing cleanup")
      #do the clean up you want
      raise StopIteration()

class MyIterable:
    def __init__(self, n):
      self.n = n
    def __iter__(self):
        return MyIterator(self.n)

for x in MyIterable(5):
  print(x)

Output:

4
9
7
5
1
doing cleanup

Note: It is better to do things here using generator function or expression, but this will make things more distinguishable.

user_35792120
Related