Pythonでディレクトリ配下の全ファイルを１行ずつ処理する
Pythonでデータを処理するコードで、
ディレクトリ配下の全ファイルに対して処理を書くとき。
- ディレクトリ配下のファイル一覧のfor文
- ファイル開くwith文
- ファイルを１行ずつ処理するwhile文
- 終了条件のif文
といった感じになるので、単純な処理の割にネストが深くなりがちです。
そこで、次のように、すっきりと書けないか考えてみました。
- ディレクトリを開くwith文
- ファイルを１行ずつ処理するfor文
# 対象のファイルが大きくない場合は、
# 全ファイルを開いてlistに入れる関数を用意すれば十分なので、
# このエントリのような面倒なことはしなくて良いです。
実装
結論から書くと、次のようなコードで実現できました。
open_dir, readlines_dir クラスを用意しておけば、__main__ 以下のように、すっきりとしたコードに出来ます。
import os
import glob
class open_dir(object):
def __init__(self, dirname):
self.files = [f for f in glob.glob('{}/*'.format(dirname)) if os.path.isfile(f)]
self.fr = None
self.fidx = -1
def __enter__(self):
self._open_nextfile()
return self
def __exit__(self, a, b, c):
if self.fr is not None:
self.fr.close()
def _open_nextfile(self):
self.fidx += 1
self.fr = open(self.files[self.fidx])
def readline(self):
ln = self.fr.readline()
if ln == '':
if self.fidx + 1 >= len(self.files):
return ''
self.fr.close()
self._open_nextfile()
return self.readline()
return ln
class readlines_dir(object):
def __init__(self, dr):
self.dr = dr
def __iter__(self):
return self
def __next__(self):
ln = self.dr.readline()
if ln == '':
raise StopIteration()
return ln.strip()
if __name__ == '__main__':
with open_dir('data') as dr:
for ln in readlines_dir(dr):
print(ln)
実装の説明
open_dir クラス
open_dirクラスは、with文を処理するクラスです。__enter__メソッドの返却値がasの後の変数に代入、__exit__メソッドの処理が、with文から抜けるときの処理です。
readlineが呼ばれる度に次の行を返し、
ファイルの末尾まで来たら、次のファイルに移るようにしています。
with文については、以下を参考にして下さい。
with文 | Python言語リファレンス
https://docs.python.org/ja/3/reference/compound_stmts.html#the-with-statement
readlines_dir クラス
readlines_dirクラスは、イテレータオブジェクトのクラスです。__next__メソッドがイテレーション結果を返却します。
イテレーションの度に、open_dirのreadlineを呼び出しています。
イテレータについては、以下を参考にして下さい。
iterator | Python 用語集
https://docs.python.org/ja/3.8/glossary.html#term-iterator
まとめ
特にデータ分析や機械学習などのコードを書いているときに、
ファイルを開くといったシステム的な処理のコードが増えると、
ビジネスロジックを追いにくくなります。
このような方法でメイン処理から、
システム的なコードを追い出しておくと、コードがわかりやすくなって良いです。
ここでは、単純にディレクトリ配下のファイルを順に処理していますが、open_dirクラスのファイルオープンの箇所を書き換えれば、
gzipで圧縮されているものを展開しながら処理したり、
サーバやクラウドストレージ上のファイルを順次ダウンロードしつつ処理したり、
と、いろいろ応用することも出来そうです。
以上。
