PythonでCSVファイルの特定の文字が含まれる行とその前後を抜き出してみる
割と行数の大きいcsvファイルがありまして、その中の要素にある文字列が含まれる行を中心にして前後に数行のデータを抜き出したいという場合があります。
わかりにくいですね。
データロガーなどで採ったデータで、1列目に時刻（またはカウンタ）、2列目に測定値（電圧とか）を記録したcsvファイルがあるとします。とても長い時間のログを取るとか、サンプリング間隔がすごく短いとかで行数が膨大になったのですが、そのなかのある時刻を中心に前後数秒（前後数カウント）だけを抜き出してグラフを作りたかったりします。
Excelでcsvファイルを開いて選択したり、テキストエディタで検索したりして編集することもできますが、面倒です。
というわけで、Pythonでやってみたいと思います。
目次
処理の流れ
csvモジュールを使ってファイル読み込み
forループとinを使って適合する要素を検索
抽出する範囲を計算
forループを使ってデータを抽出
出力
コード
下記のコードでtrim_csv_row.pyというファイルにしました。
import csv
import sys
arguments = sys.argv
fn = ''
search_string = ''
prev_range = 0
after_range = 0
search_column = 0
fn_out = ''
# コマンドラインオプション
for i in arguments:
if '-col' in i:
search_column = int(arguments[arguments.index('-col') + 1])
continue
if '-f' in i:
fn = arguments[arguments.index('-f') + 1]
continue
if '-o' in i:
fn_out = arguments[arguments.index('-o') + 1]
continue
if '-s' in i:
search_string = arguments[arguments.index('-s') + 1]
continue
if '-b' in i:
prev_range = int(arguments[arguments.index('-b') + 1])
continue
if '-a' in i:
after_range = int(arguments[arguments.index('-a') + 1])
continue
if fn == '':
print('No input file.')
exit()
if search_string == '':
print('No search string.')
exit()
# ファイル読み込み
with open(fn, mode='r', newline='') as f_in:
reader = csv.reader(f_in)
data_array = [row for row in reader]
# 基準となる行の検索
index_position = None
for row_counter in range(len(data_array)):
if search_string in data_array[row_counter][search_column]:
index_position = row_counter
break
if index_position == None:
print('Not matched.')
exit()
# 出力範囲の計算
pickup_start = index_position - prev_range
pickup_end = index_position + after_range
if pickup_start < 0:
pickup_end = pickup_end + abs(pickup_start)
pickup_start = 0
if pickup_end > len(data_array) - 1:
pickup_start = pickup_start - (pickup_end-len(data_array) + 1)
if pickup_start < 0:
pickup_start = 0
pickup_end = len(data_array) - 1
# 出力用リストの生成
output_array = []
for row_counter in range(pickup_start, pickup_end+1):
output_array.append(data_array[row_counter])
if fn_out == '':
for i in output_array:
print(i)
else:
with open(fn_out, mode = 'wt', newline='') as f_out:
csvout = csv.writer(f_out)
csvout.writerows(output_array)
コマンドラインオプションで、いろいろ指定できるようにしてみました。
-f : 入力するファイル名
-o : 出力するファイル名（省略した場合は標準出力）
-s : 検索文字列
-col : 検索する列
-b : 該当する行よりこの数だけ前の行を出力する
-a : 該当する行よりこの数だけ後の行を出力する
出力する行の指定が元のデータの範囲を超えたら、できるだけ同じ幅の範囲を出力するように調整します。 グラフとかにしたときに、幅が同じ方が扱いやすいですからね。
テスト用のデータ
データはこんな感じにしてみました。
ab,123
cd,456
ef,789
gh,123
ij,456
kl,789
mn,123
op,456
qr,789
st,123
uv,456
wx,789
yz,0
試してみた
> python trim_csv_row.py -f data.csv -s f -col 0 -b 1 -a 2 -o result.csv
result.csvの中身はこうなりました。
cd,456
ef,789
gh,123
ij,456
一応指定通り出力できてますな。
公開日
広告