A cut(1)-style tool for working with CSV files

csv.py

#!/usr/bin/env python
from __future__ import print_function
import argparse
import csv
import os
import sys


def fields(string):
    parsed_fields = []
    field_strings = string.split(',')
    for field in field_strings:
        if '-' in field:
            (start, end) = map(int, field.split('-'))
            parsed_fields.extend(range(start, end + 1))
        else:
            parsed_fields.append(int(field))
    return parsed_fields
 
 
def main(args):
    exit_code = os.EX_OK
 
    infiles = args.file
    selected_fields = args.fields
    delimiter = args.output_delimiter
 
    if delimiter is None:
        output = csv.writer(sys.stdout)
    else:
        output = sys.stdout
 
    for infile in infiles:
        if infile == '-':
            infile = sys.stdin
        else:
            infile = open(infile, 'rb')
        reader = csv.reader(infile)
        for row in reader:
            out = []
            if selected_fields is None:
                out = row
            else:
                for field in selected_fields:
                    try:
                        out.append(row[field - 1])
                    except IndexError:
                        pass
            if delimiter is None:
                output.writerow(out)
            else:
                print(delimiter.join(out), file=output)
            sys.stdout.flush()
 
    return exit_code
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--fields', '-f', type=fields)
    parser.add_argument('--output-delimiter', '-o', type=str)
 
    parser.add_argument('file', nargs="*", default=['-'])
 
    args = parser.parse_args()
    try:
        exit_code = main(args)
    except IOError as ioe:
        if ioe.errno == 32:  # broken pipe
            pass
        else:
            raise
    except KeyboardInterrupt:
        pass