Utility to selectively output CSV files

csv

#!/usr/bin/env python

import csv
import sys

if len(sys.argv) < 2:
    print("usage: csv <filename> <opt: col1,col2,...>")
    sys.exit(1)

filename = sys.argv[1]

def get_columns(f):
    firstline = f.readline().strip()
    return csv.reader([firstline]).next()

with open(filename) as f:
    # Print specified columns from each row
    if len(sys.argv) == 3:
        columns_mapping = {}
        for idx, col in enumerate(get_columns(f)):
            columns_mapping[col] = idx
        print_columns = map(lambda c: columns_mapping[c], sys.argv[2].strip().split(','))
        reader = csv.reader(f)
        for row in reader:
            to_print = '\t'.join(map(lambda c: row[c], print_columns))
            print to_print
    # If no specific column to print, print all columns
    else:
        for idx, col in enumerate(get_columns(f)):
            print(str(idx) + ': ' + col)