如何使用列表标识键来打印未知数量的字典键？

Question

So I have searched and found certain things that have helped me put this code together, but I haven't had any luck with this last part. What I want to do is read in a file that is comma separated, space separated, or tab separated, set the headers as keys and the data as the values, and then write only certain columns (number of columns is unknown) to an output file. Example.txt is shown below:

col1, col2, col3, col4, col5
1, 11, 21, 31, 41
2, 12, 22, 32, 42
3, 13, 23, 33, 43
4, 14, 24, 34, 44

So far, here is the working code that I have so far.

import csv
import sys

file = sys.argv[1] # name of file is example.txt
columns = sys.argv[2:] # order: col1, col3, col5

with open(file, 'r') as csvfile:
    with open('table.out', 'w') as file_out:
        file.out_write(columns[0] + '\t' + columns[1] + '\t' + columns[2] + '\n')

        reader = csv.DictReader(csvfile)
        for row in reader:
            file_out.write(row[columns[0]] + '\t' + row[columns[1]] + '\t' + row[columns[2]] + '\n') 

Results:

col_1    col_3    col_5
1    21    41
2    22    42
3    23    43
4    24    44

This code works great if the number of columns were a fixed number, but the number of columns to be written can vary. For example, sometimes I may want to only grab col1, col2 and other times I may want to grab col2, col3, col4, col5 in no particular order.

So my question is, how can I modify the above code such that I can have any number of columns be written to the output file using dictionaries in Python 3.X?

Answer 1

You can adapt this to your needs, but basically using the join function will be really helpful + a list comprehension.

import csv
import sys

file = sys.argv[1]
columns = sys.argv[2:]

with open(file) as f:
    myread = csv.DictReader(f)
    for row in myread:
        print('\t'.join([row[i] for i in columns]))

Answer 2

import csv
import sys

file = sys.argv[1] # name of file is example.txt
columns = sys.argv[2:] # order: col1, col3, col5
n_columns=len(columns)

with open(file, 'r') as csvfile:
    with open('table.out', 'w') as file_out:
        for i in range(0,n_columns):
            file_out.write(columns[i] + '\t')
        file_out.write('\n')

        reader = csv.DictReader(csvfile)
        for row in reader:
            for i in range(0,n_columns):
                file_out.write(row[columns[i]] + '\t')
            file_out.write('\n')

so, i've modified your code a little bit. to write a variable number of columns you can use a for statement that go from 0 to the length of the columns list.