-
Notifications
You must be signed in to change notification settings - Fork 1
/
dict_merge.py
executable file
·52 lines (41 loc) · 1.34 KB
/
dict_merge.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import sys
import csv
import smart_csv
def read_dict(fname, key, vals):
d = {}
with open(fname) as fin:
c = smart_csv.csv_open(fin)[0]
headers = c.next()
keyidx = headers.index(key)
validx = [headers.index(val) for val in vals]
for x in c:
d[x[keyidx]] = [x[i] for i in validx]
return d
def merge_csv(orig_csv, target_idx, d, key, vals, out_csv, keep):
with open(orig_csv) as fin:
cin = smart_csv.csv_open(fin)[0]
with open(out_csv, 'wt') as fout:
cout = csv.writer(fout)
headers = cin.next()
keyidx = headers.index(key)
def merge_row(row, vals):
del row[keyidx]
return row[:target_idx] + vals + row[target_idx:]
cout.writerow(merge_row(headers, vals))
for x in cin:
cout.writerow(merge_row(x, d.get(x[keyidx], [''] * len(vals))))
if __name__ == '__main__':
orig_csv = sys.argv[1]
target_idx = int(sys.argv[2])
out_csv = sys.argv[3]
dict_csv = sys.argv[4]
if sys.argv[5] == 'keep':
keep = 1
elif sys.argv[5] == 'kill':
keep = 0
else:
raise ValueError
key = sys.argv[6]
vals = sys.argv[7:]
d = read_dict(dict_csv, key, vals)
merge_csv(orig_csv, target_idx, d, key, vals, out_csv, keep)