CODEWITHSUNDEEP
pythonfile-ioiterator
jsj
jsj

Reputation: 9421

Python Custom Iterator: Close a file on StopIteration

I have written an iterator class that opens a file in it's __init__.

def __init__(self, path):
    self.file = open(path, "r")

How do I close that file automatically when the iteration is finished?

Complete class:

class Parse(object):
    """A generator that iterates through a CC-CEDICT formatted file, returning
   a tuple of parsed results (Traditional, Simplified, Pinyin, English)"""
    def __init__(self, path):
        self.file = open(path, "r")

    def __iter__(self):
        return self

    def __is_comment(self, line):
        return line.startswith("#")

    def next(self):
        #This block ignores comments.
        line = self.file.readline()
        while line and self.__is_comment(line):
            line = self.file.readline()

        if line:
            working = line.rstrip().split(" ")
            trad, simp = working[0], working[1]
            working = " ".join(working[2:]).split("]")
            pinyin = working[0][1:]
            english = working[1][1:]
            return trad, simp, pinyin, english

        else:
            raise StopIteration()

Upvotes: 9

Views: 5493

Answers (2)

Ned Batchelder
Ned Batchelder

Reputation: 375814

A better way to write the whole thing would be to keep the opening and the iteration in one place:

class Parse(object):
    """A generator that iterates through a CC-CEDICT formatted file, returning
    a tuple of parsed results (Traditional, Simplified, Pinyin, English)"""
    def __init__(self, path):
        self.path = path

    def __is_comment(self, line):
        return line.startswith("#")

    def __iter__(self):
        with open(self.path) as f:
            for line in f:
                if self.__is_comment(line):
                    continue

                working = line.rstrip().split(" ")
                trad, simp = working[0], working[1]
                working = " ".join(working[2:]).split("]")
                pinyin = working[0][1:]
                english = working[1][1:]
                yield trad, simp, pinyin, english

This will wait to open the file until you really need it, and will automatically close it when done. It's also less code.

If you really want to get into the "generators are awesome!" mindset:

def skip_comments(f):
    for line in f:
        if not.startswith('#'):
            yield line

...

    def __iter__(self):
        with open(self.path) as f:
            for line in skip_comments(f):
                working = ....

Upvotes: 16

Martijn Pieters
Martijn Pieters

Reputation: 1123790

You need to explicitly close it as soon as StopIteration is raised. In this case, simply call .close() when you raise StopIteration yourself.

def next(self):
    #This block ignores comments.
    line = self.file.readline()
    while line and self.__is_comment(line):
        line = self.file.readline()

    if line:
        working = line.rstrip().split(" ")
        trad, simp = working[0], working[1]
        working = " ".join(working[2:]).split("]")
        pinyin = working[0][1:]
        english = working[1][1:]
        return trad, simp, pinyin, english

    else:
        self.file.close()
        raise StopIteration()

Since no other code in your .next() method could trigger a StopIteration this suffices.

If you did use next() on another iterator inside your own .next() you'd have to catch StopIteration with an except StopIteration: handler and reraise the exception.

This only handles the StopIteration case. If you want to handle other situations (not exhausting the iterator) you'll need to handle that situation separately. Making your class a Context Manager as well could help with that. Users of your iterator would then use the object in a with statement before iterating over it, and when the with suite is exited the file could be closed regardless. You may want to mark your iterator as 'done' as well in that case:

_closed = False

def next(self):
    if self._closed:
        raise StopIteration

    line = self.file.readline()
    while line and self.__is_comment(line):
        line = self.file.readline()

    if line:
        working = line.rstrip().split(" ")
        trad, simp = working[0], working[1]
        working = " ".join(working[2:]).split("]")
        pinyin = working[0][1:]
        english = working[1][1:]
        return trad, simp, pinyin, english

    else:
        self.file.close()
        self._closed = True
        raise StopIteration()  

def __enter__(self):
    return self

def __exit__(self, type_, value, tb):
    self.file.close()  # multiple calls to .close() are fine
    self._closed = True

Upvotes: 1

Related Questions