I needed to make a small script to read a csv file and store 2 columns in Redis. Can you guys take a look over the code and let me know how I can improve my code style?
#!/usr/bin/env/ python
"""Proces csv file and store key-value pair of columns in redis"""
import csv, redis, json
from sys import argv, exit
REDIS_HOST = 'localhost'
conn = redis.Redis(REDIS_HOST)
def get_csv_data(csv_file, **columns):
col1, col2 = columns['columns']
with open(csv_file, encoding='utf-8') as csvf:
csv_data = csv.reader(csvf)
return [(i[col1], i[col2]) for i in csv_data]
def store_data(data):
for i in data:
if i[1] not in conn:
conn.set(i[0], i[1])
return data
def main():
if len(argv) < 2:
exit("usage: python {file} file.csv [column1, column2]".format(file=__file__))
csv_file = argv[1]
columns = list(map(int,argv[2:4])) or [1, 3]
data = get_csv_data(csv_file, columns=columns)
print (json.dumps(store_data(data)))
if __name__ == '__main__':
main()