CODEWITHSUNDEEP
pythoncontextmanager
wwii
wwii

Reputation: 23773

Context Manager Type that iterates over multiple files - testing

Id like to use itertools.izip() to iterate over the lines of multiple files. I have made a context manager type to ensure all the files will be closed when used with with. This seems to work:

class Files_Iterator(object):
    """Sequential line iteration from multiple files
    """
    def __init__(self, files):
        """files --> list or tuple of str, file paths
        """
        self.files = files
    def __enter__(self):
        print 'opening files'
        self.files = map(open, self.files)
        return it.izip(*self.files)
    def __exit__(self, *exception_args):
        print 'closing files'
        for arg in exception_args:
            print arg,
        for thing in self.files:
            thing.close()
        return False

Two questions:

  1. Have I implemented this correctly?
  2. Can I test this to ensure the files are getting closed or do I just trust it?

I used print statements to signal when __exit__ is called - is that a sufficient test?

>>> with Files_Iterator(['file1.txt', 'file2.txt']) as files:
    for lines in files:
        print lines
        raise Exception


opening files
('File1Line1\n', 'File2Line1\n')
closing files
<type 'exceptions.Exception'>  <traceback object at 0x0313DFD0>

Traceback (most recent call last):
  File "<pyshell#48>", line 4, in <module>
    raise Exception
Exception
>>>

Upvotes: 1

Views: 403

Answers (1)

Aaron Hall
Aaron Hall

Reputation: 395523

It looks fine, yes you can trust it, but I would explicitly name the arguments to __exit__:

def __exit__(self, exc_type, exc_value, traceback):
    print 'closing files'
    for arg in (exc_type, exc_value, traceback):
        print arg,
    for f in self.files:
        f.close()
    return False # does not suppress the exception.

And when the function exits, if there is an exception, it will be processed normally.

Upvotes: 1

Related Questions