Ignore:
Timestamp:
Feb 17, 2009, 4:51:44 PM (16 years ago)
Author:
rwilson
Message:

Add the 'be_green' option to add_csv_header().

File:
1 edited

Legend:

Unmodified
Added
Removed
  • anuga_core/source/anuga/lib/add_csv_header/add_csv_header.py

    r6355 r6357  
    4545    # copy header to output file, then input file
    4646    tmp_fd.write(header_string)
    47     data = fd.readlines()
    48     columns = data[0].strip().split(',')
    4947
    50     if len(columns) != len(header_list):
    51         msg = ("add_csv_header: File %s has %d columns but header "
    52                "has %d columns" % (file, len(columns), len(header_list)))
    53         raise Exception, msg
     48    if be_green:
     49        first_line = True
     50        for line in fd:
     51            if first_line:
     52                first_line = False
     53                columns = line.strip().split(',')
     54                if len(columns) != len(header_list):
     55                    msg = ("add_csv_header: File %s has %d columns but header "
     56                           "has %d columns" % (file, len(columns), len(header_list)))
     57                    raise Exception, msg
     58            tmp_fd.write(line)
     59    else:
     60        data = fd.readlines()
     61        columns = data[0].strip().split(',')
    5462
    55     data = ''.join(data)
    56     tmp_fd.write(data)
     63        if len(columns) != len(header_list):
     64            msg = ("add_csv_header: File %s has %d columns but header "
     65                   "has %d columns" % (file, len(columns), len(header_list)))
     66            raise Exception, msg
     67
     68        data = ''.join(data)
     69        tmp_fd.write(data)
    5770
    5871    # close and rename all files
     
    6477if __name__ == '__main__':
    6578    import sys
     79    import os
    6680
    6781    file_data = '1,2,3\n4,5,6\n7,8,9'
    6882    header = ['alpha', 'bravo', 'charlie']
    6983    filename = '/tmp/add_csv_header.csv'
     84    filename2 = '/tmp/add_csv_header2.csv'
     85
     86######
     87# Create file and test function.
     88######
    7089
    7190    # create test file
     
    89108    assert expected == data, msg
    90109
     110######
     111# Test the 'be_green' option.
     112######
     113
     114    # create test file
     115    fd = open(filename2, 'w')
     116    fd.write(file_data)
     117    fd.close()
     118
     119    add_csv_header(filename2, header, be_green=True)
     120
     121    # read test file
     122    fd = open(filename2, 'r')
     123    data = fd.readlines()
     124    fd.close
     125
     126    # check if data as expected
     127    data = ''.join(data)
     128    header_string = ','.join(header) + '\n'
     129    expected = header_string + file_data
     130
     131    msg = 'Expected data:\n%s\ngot:\n%s' % (expected, data)
     132    assert expected == data, msg
     133
     134    os.remove(filename)
     135    os.remove(filename2)
Note: See TracChangeset for help on using the changeset viewer.