Changeset 6357 for anuga_core/source/anuga/lib/add_csv_header/add_csv_header.py

Timestamp:

Feb 17, 2009, 4:51:44 PM (15 years ago)

Author:

rwilson

Message:

Add the 'be_green' option to add_csv_header().

File:

• anuga_core/source/anuga/lib/add_csv_header/add_csv_header.py (modified) (3 diffs)

anuga_core/source/anuga/lib/add_csv_header/add_csv_header.py

@@ -45 +45 @@
     # copy header to output file, then input file
     tmp_fd.write(header_string)
-    data = fd.readlines()
-    columns = data[0].strip().split(',')
 
-    if len(columns) != len(header_list):
-        msg = ("add_csv_header: File %s has %d columns but header "
-               "has %d columns" % (file, len(columns), len(header_list)))
-        raise Exception, msg
+    if be_green:
+        first_line = True
+        for line in fd:
+            if first_line:
+                first_line = False
+                columns = line.strip().split(',')
+                if len(columns) != len(header_list):
+                    msg = ("add_csv_header: File %s has %d columns but header "
+                           "has %d columns" % (file, len(columns), len(header_list)))
+                    raise Exception, msg
+            tmp_fd.write(line)
+    else:
+        data = fd.readlines()
+        columns = data[0].strip().split(',')
 
-    data = ''.join(data)
-    tmp_fd.write(data)
+        if len(columns) != len(header_list):
+            msg = ("add_csv_header: File %s has %d columns but header "
+                   "has %d columns" % (file, len(columns), len(header_list)))
+            raise Exception, msg
+
+        data = ''.join(data)
+        tmp_fd.write(data)
 
     # close and rename all files
@@ -64 +77 @@
 if __name__ == '__main__':
     import sys
+    import os
 
     file_data = '1,2,3\n4,5,6\n7,8,9'
     header = ['alpha', 'bravo', 'charlie']
     filename = '/tmp/add_csv_header.csv'
+    filename2 = '/tmp/add_csv_header2.csv'
+
+######
+# Create file and test function.
+######
 
     # create test file
@@ -89 +108 @@
     assert expected == data, msg
 
+######
+# Test the 'be_green' option.
+######
+
+    # create test file
+    fd = open(filename2, 'w')
+    fd.write(file_data)
+    fd.close()
+
+    add_csv_header(filename2, header, be_green=True)
+
+    # read test file
+    fd = open(filename2, 'r')
+    data = fd.readlines()
+    fd.close
+
+    # check if data as expected
+    data = ''.join(data)
+    header_string = ','.join(header) + '\n'
+    expected = header_string + file_data
+
+    msg = 'Expected data:\n%s\ngot:\n%s' % (expected, data)
+    assert expected == data, msg
+
+    os.remove(filename)
+    os.remove(filename2)