csvデータの加工

Python

いつものように備忘録として。。。

csvファイルの読み込み

python SCRIPTFILE csv_file elements

という感じで使うことを前提。elementsは抽出したい要素を指定(複数可)
csvファイルの書式は下の通り

ID,SEX,AGE,Pref,...
1,0,33,25,...
2,1,24,10,...
...

まずはファイルを読み込んで内容を返す関数を作成

def readFile(filename):
	print 'reading ',filename,'...'
	read_file = open(filename,'r')
	reader = csv.reader(read_file)
	list = []
	for line in reader:
		list.append(line)
	return list

まずハマったのは、csv.readerで取得したオブジェクト？からfor文で配列に入れるところ。
ファイルに書きこんで実行。じゃなくてインタプリタで試行錯誤してたからハマった問題やな。。
一度for line in reader:を回して処理をしてしまうとreaderの中身は空になってしまう。
print readerでオブジェクトの存在は確認できるのに、reader=csv(ryとやらずにfor分からやろうとするとリストにうまくデータを格納できなかった。

csvファイルへ書き込み

↑と同様に関数を作成

def writeFile(write_datum,filename):
	print 'writing ',filename,'.csv...'
	write_file = file(filename + ".csv","w")
	writer = csv.writer(write_file,lineterminator="\n")
	i = 0
	while(i < len(write_datum)):
		writer.writerow(write_datum[i])
		i = i + 1

write_datumの中身は↓のとおり

[ ['AGE', 'MARRIAGE'], ['40', '1'], ['49', '2'], ['45', '1'],...]

リストへの追加方法はappendとextendを使う方法があるんだが、違いが分からなかったので試してみた。

その前にcsvから任意の行を抽出する関数を作成

def extractColumns(list,elements):
	#要素が存在しなければ中止
	write_datum = []
	rows = []
	datum = []
	#抽出する要素が存在するかチェック、存在するならリストに追加
	for element in elements:
		if not element in list[0]:
			print element,'is not exist!'
			quit()
		n = list[0].index(element)
		rows.append(n)
	#１行ごとにwrite_datumに追加
	#["list[i][row1]","list[i][row2]"...]という感じ
	#csvファイルでは list[i][row1],list[i][row2]...というようになる
	i = 0
	while(i < len(list)):
		datum = []
		for row in rows:
			datum.append(list[i][row])
		write_datum.append(datum)
		i = i + 1
	return write_datum

datum.append→datum.extendに変えると↓のように変わる。

[ ['A', 'G', 'E', 'M', 'A', 'R', 'R', 'I', 'A', 'G', 'E'], ['4', '0', '1'], ...]

write_datum.append→write_datum.extendに変えると↓のように変わる。

['AGE', 'MARRIAGE', '40', '1', '49', '2', '45',...]

datum.append→datum.extend、write_datum.append→write_datum.extendに変えると↓のように変わる。

['A', 'G', 'E', 'M', 'A', 'R', 'R', 'I', 'A', 'G', 'E', '4', '0', '1', ...]

ソースコードは続きに乗せておくので興味がある人はどうぞ(^_^)つ旦

# -*- coding: utf-8 -*-
#for python2.6
#usage: python process.py csv_file elements

import csv
import sys

############
#          #
# 関数定義 #
#          #
############


#filenameを読み込んで内容をリストとして返す
def readFile(filename):
	print 'reading ',filename,'...'
	read_file = open(filename,'r')
	reader = csv.reader(read_file)
	list = []
	for line in reader:
		list.append(line)
	return list


#write_datum[]をfilename.csvとして書きこむ
def writeFile(write_datum,filename):
	print 'writing ',filename,'.csv...'
	write_file = file(filename + ".csv","w")
	writer = csv.writer(write_file,lineterminator="\n")

	i = 0
	while(i < len(write_datum)):
		writer.writerow(write_datum[i])
		i = i + 1


#引数を取得、読み込むファイル、抽出する行を設定
def conf():
	arg_list = sys.argv
	#引数が足りなければ通知
	if len(arg_list) < 3:
		print 'USAGE: python process.py csv_file elements'
		quit()
	return arg_list


#csvから任意の行を抽出、elementsは配列
def extractColumns(list,elements):
	#要素が存在しなければ中止
	write_datum = []
	rows = []
	datum = []
	#抽出する要素が存在するかチェック、存在するならリストに追加
	for element in elements:
		if not element in list[0]:
			print element,'is not exist!'
			quit()
		n = list[0].index(element)
		rows.append(n)
	#１行ごとにwrite_datumに追加
	#["list[i][row1]","list[i][row2]"...]という感じ
	#csvファイルでは list[i][row1],list[i][row2]...というようになる
	i = 0
	while(i < len(list)):
		datum = []
		for row in rows:
			datum.extend(list[i][row])
		write_datum.append(datum)
		i = i + 1
	return write_datum


########
#      #
# 処理 #
#      #
########


args = conf()
#抽出する要素を配列で取得
elements = args[2:]

#csvファイル読み込み
readfile = args[1]
list = readFile(readfile)

write_datum = extractColumns(list,elements)

#filenameの作成
writefile = readfile[:-4]
for x in elements:
	writefile += '_' + x[:3]
writeFile(write_datum,writefile)


print 'finished processing!'