将过滤后的CSV文件写入新文件并遍历文件夹

Question

I have been trying initially to create a program to go through one file and select certain columns that will then be moved to a new text file. So far I have

    import os, sys, csv
    os.chdir("C://Users//nelsonj//Desktop//Master_Project")
    with open('CHS_2009_test.txt', "rb") as sitefile:
    reader = csv.reader(sitefile, delimiter=',')
    pref_cols = [0,1,2,4,6,8,10,12,14,18,20,22,24,26,30,34,36,40]

    for row in reader:
        new_cols = list(row[i] for i in pref_cols)
        print new_cols

I have been trying to use the csv functions to write the new file but I am continuosly getting errors. I will eventually need to do this over a folder of files, but thought I would try to do it on one before tackling that.

Code I attempted to use to write this data to a new file

    for row in reader:
        with open("CHS_2009_edit.txt", 'w') as file:
            new_cols = list(row[i] for i in pref_cols)
            newfile = csv.writer(file)
            newfile.writerows(new_cols)

This kind of works in that I get a new file, but in only prints the second row of values from my csv, ie, not the header values and places commas in between each individual character, not just copying over the original columns as they were.

I am using PythonWin with Python 2.6(from ArcGIS)

Thanks for the help!

NEW UPDATED CODE

   import os, sys, csv

   path = ('C://Users//nelsonj//Desktop//Master_Project')

   for filename in os.listdir(path):

       pref_cols = [0,1,2,4,6,8,10,12,14,18,20,22,24,26,30,34,36,40]
       with open(filename, "rb") as sitefile:
           with open(filename.rsplit('.',1)[0] + "_Master.txt", 'w') as output_file:
               reader = csv.reader(sitefile, delimiter=',')
               writer = csv.writer(output_file)
               for row in reader:
                   new_row = list(row[i] for i in pref_cols)
                   writer.writerow(new_row)
                   print new_row

Getting list index out of range for the new_row, but it seems to still be processing the file. Only thing I can't get it to do now is loop through all files in my directory. Here's a hyperlink to Screenshot of data text file

Answer 1

Try this:

 new_header = list(row[i] for i in pref_cols if i in row)

That should avoid the error, but it may not avoid the underlying problem. Would you paste your CSV file somewhere that I can access, and I'll fix this for you?

Answer 2

For your purpose of filtering, you don't have to treat the header differently from the rest of the data. You can go ahead remove the following block:

    headers = reader.next()
    for row in headers:
        new_header = list(row[i] for i in pref_cols)
        print new_header  

Your code did not work because you treated headers as a list of rows, but headers is just one row.

Update

This update deals with writing the CSV data to a new file. You should move the open statement above the for row...

with open("CHS_2009_edit.txt", 'w') as output_file:
    writer = csv.writer(output_file)
    for row in reader:
        new_cols = list(row[i] for i in pref_cols)
        writer.writerows(new_cols)

Update 2

This update deals with the header output problem. If you followed my suggestions, you should not have this problem. I don't know what your current code looks like, but it looks like you supplies a string where the code expects a list. Here is the code that I tried on my system (using my made-up data) and it seems to work:

pref_cols = [...] # <<=== Should be set before entering the loop
with open('CHS_2009_test.txt', "rb") as sitefile:
    with open('CHS_2009_edit.txt', 'w') as output_file:
        reader = csv.reader(sitefile, delimiter=',')
        writer = csv.writer(output_file)
        for row in reader:
            new_row = list(row[i] for i in pref_cols)
            writer.writerow(new_row)

One thing to notice: I use writerow() to write a single row, where you use writerows() -- that makes a difference.