Search code examples
pythoncsvfor-loopcontextmanager

how to write results in a for loop to multiple CSV files


I have

with open ('~/abc.csv', 'w') as f:
    write1 = csv.write(f)
    write1.writerow(['header1', 'header2', 'header3', 'header4'])

with open ('~/def.csv', 'w') as g:
    write2 = csv.write(g)
    write2.writerow(['header1', 'header2', 'header3', 'header4', 'header5', 'header6'])

for iteration in a_list:
    perform calculations
    result1 = ([h1, h2, h3, h4],[l1, l2, l3, l4],[m1, m2, m3, m4], ...,[])
    for pa in result1:
        write1.writerow(pa)

    def fun(result1):
        result2 = ([n1, n2, n3, n4, n5, n6],[p1, p2, p3, p4, p5, p6], [], ...[])
        for pb in result2:
             write2.writerow(pb)

Expecting two csv files as

'header1', 'header2', 'header3', 'header4'
h1, h2, h3, h4
l1, l2, l3, l4
m1, m2, m3, m4
      :

and

'header1', 'header2', 'header3', 'header4', 'header5', 'header6'
n1, n2, n3, n4, n5, n6
p1, p2, p3, p4, p5, p6

This can be easily done when all the iterations are completed and separate appended(list) result1 and result2 can be easily written to separate files using writer.writerows(pa). However, I would like to write to csv files in each iteration, so that I wouldn't miss if I have to break the loop for some reason.

Thank you!


Solution

  • You can place two files in a context manager and have two csv.writer objects (since Python 2.7):

    with open ('~/abc.csv', 'w') as f, open ('~/def.csv', 'w') as g:
        writer1 = csv.writer(f)
        writer2 = csv.writer(g)
    
        writer1.writerow(['header1', 'header2', 'header3', 'header4'])    
        writer2.writerow(['header1', 'header2', 'header3', 'header4', 'header5', 'header6'])
    
        for iteration in a_list:
            # perform calculations
            for pa in result1:
                writer1.writerow(pa)
            for pb in result2:
                writer2.writerow(pb)
    

    From the docs (which are the same for Python 2.7 up to 3.5).

    With more than one item, the context managers are processed as if multiple with statements were nested:

    with A() as a, B() as b:
        suite
    

    is equivalent to

    with A() as a:
        with B() as b:
            suite
    

    This suggests the other way of doing it which is to have the nesting go one layer deeper.

    with open ('~/abc.csv', 'w') as f:
        with open ('~/def.csv', 'w') as g:
            writer1 = csv.writer(f)
            writer2 = csv.writer(g)
            ...