In Python, how do I iterate over one iterator and then another?

I'd like to iterate two different iterators, something like this:

file1 = open('file1', 'r')
file2 = open('file2', 'r')
for item in one_then_another(file1, file2):
    print item

Which I'd expect to print all the lines of file1, then all the lines of file2.

I'd like something generic, as the iterators might not be files, this is just an example. I know I could do this with:

for item in [file1]+[file2]:

but this reads both files into memory, which I'd prefer to avoid.

Use itertools.chain:

from itertools import chain
for line in chain(file1, file2):
   pass

fileinput module also provides a similar feature:

import fileinput
for line in fileinput.input(['file1', 'file2']):
   pass

You can also do it with simple generator expression:

for line in (l for f in (file1, file2) for l in f):
    # do something with line

with this method you can specify some condition in expression itself:

for line in (l for f in (file1, file2) for l in f if 'text' in l):
    # do something with line which contains 'text'

The example above is equivalent to this generator with loop:

def genlinewithtext(*files):
    for file in files:
        for line in file:
            if 'text' in line:
                yield line

for line in genlinewithtext(file1, file2):
    # do something with line which contains 'text'

I think the most Pythonic approach to this particular file problem is to use the fileinput module (since you either need complex context managers or error handling with open), I'm going to start with Ashwini's example, but add a few things. The first is that it's better to open with the U flag for Universal Newlines support (assuming your Python is compiled with it, and most are), (r is default mode, but explicit is better than implicit). If you're working with other people, it's best to support them giving you files in any format.

import fileinput

for line in fileinput.input(['file1', 'file2'], mode='rU'):
   pass

This is also usable on the command line as it will take sys.argv[1:] if you do this:

import fileinput

for line in fileinput.input(mode='rU'):
   pass

And you would pass the files in your shell like this:

$ python myscript.py file1 file2

In Python, how do I iterate over one iterator and then another?

Related

Recent Posts