Python script to split columns from csv file

gistfile1.py

#!/usr/bin/python
# Input:  CSV file with headers, column names
#   x,y,z
#   1,2,3
#   4,5,6
# Output: CSV file
# Usage:   
#  ./split-columns.py [IN_CSV] [COLUMN NAMES]
# Example: ./split-columns.py in.csv x z
import sys

def main(args):
    in_fn = args[0]
    with open(in_fn, 'r') as f:
        lines = f.readlines()
    header = lines[0].strip().split(',')
    
    idx_list = [header.index(a) for a in args[1:]]
    b = map(lambda x: [x.strip().split(',')[i] for i in idx_list], lines)
    
    sys.stdout.write('\n'.join([','.join(a) for a in b]))
    sys.stdout.write('\n')
    
if __name__ == '__main__':
    main(sys.argv[1:])