How to fill new columns in a csv file through command line

Question

I have this text file:-

User_ID,Date,Num_1,Num_2,Com_ID
101,2015-04-13,12,21,1011
102,2014-04-03,1,7,1002

I am doing some basic manipulation in the input file and writing it back into the sample_output.txt file.

#!/bin/usr/python
import datetime
import csv
import sys
file_name='sample.txt'
with open(file_name,'rb') as f:               
    reader = csv.reader(f)  
    headers = reader.next()
    p=[]
    for row in reader:            
        row[0] = row[0].zfill(6) 
        row[2] = row[2].zfill(6)
        row[3] = row[3].zfill(6)
        row[4] = row[4].zfill(6)
        row[1] = row[1][5:7] + "-" + row[1][8:10] + "-" + row[1][:4]
        p.append(row[:5])
with open('sample_out.txt', 'wb') as ofile:     
        header = ['User_ID','Date','Num_1','Num_2','Com_ID','Dept','Location']
        writer = csv.DictWriter(ofile, fieldnames=header)
        writer.writeheader()
        col_fill = ''
        writer.writerows({col:row_item for row_item,col in zip(row+[col_fill]*n,header)} for row in p)

I need to add two new columns into the output file(sample_out.txt) and pass the values of those two columns via command line such that when I run the script:- python script.py BUS SAF the output should be :-

User_ID,Date,Num_1,Num_2,Com_ID,Dept,Location
000101,04-13-2015,000012,000021,001011,BUS,SAF
000102,04-03-2014,000001,000007,001002,BUS,SAF

I can create the new columns.How do I fill the new columns by passing values via command line.

Answer 1

this is very similar to your other post which I also answered... here I instead use the argparse library to give you input options to use like:

$ python pyscript.py -cl Dept Location -vl BUS SAF

you can also use like:

$ python pyscript.py --column_names Dept Location --fill_values BUS SAF

here is the code:

import argparse

# parse arguments from commandline if available
parser = argparse.ArgumentParser()
parser.add_argument('-cl', '--column_names', default=None, nargs='+', type=str)
parser.add_argument('-vl', '--fill_values', default=None,  nargs='+', type=str)

extra_headers = parser.parse_args().column_names
col_fill = parser.parse_args().fill_values

# make sure col_fill is the same length as extra_headers    
if len(col_fill) == len(extra_headers):
    raise ValueError('Wrong number of column fills for number of headers!')
header.extend(extra_headers)

writer = csv.DictWriter(outcsv, fieldnames = header)
writer.writeheader()

# extend p with two columns of blank data
writer.writerows({col: row_item for row_item,col in zip(row+col_fill,header)} for row in p)