Cucco’s Compute Hack

コンピュータ関係の記事を書いていきます。

複数ファイルを連続して読むプログラム

複数に分割されたファイルを決まった順番で、連続して読む。
csv.readerっぽく。forで1行ずつ読めるようにイテレータとして実装。

import csv

class MyIterator(object):
    def __init__(self, readFileNames, skipHeader=False):
        self.readFileNames = readFileNames
        self.currentFileIndex = -1
        self.fp = None
        self.csv_reader = None
        self.skipHeader = skipHeader

    def __iter__(self):
        # next()はselfが実装してるのでそのままselfを返す
        return self

    def __next__(self):

        if self.csv_reader is None:
            #1個目のファイルを開く
            self.nextfile()

        try:
            value=next(self.csv_reader)
            return value
        except StopIteration:
            #終端まで来たら、このStopIterationが投げられる。
            try:
                self.nextfile()
                value=next(self.csv_reader)
                return value
            except StopIteration:
                # 次のファイルがなかったので、例外を投げてイテレータを終わる。
                raise (StopIteration)
        except:
            # ここには来ない。
            raise (StopIteration)
    
    def nextfile(self):
        self.currentFileIndex = self.currentFileIndex + 1
        if self.fp is not None:
            self.fp.close()
        if self.currentFileIndex > len(self.readFileNames)-1:
            # 次のファイルはないので例外を投げる
            raise(StopIteration)
        else:
            self.fp=open(self.readFileNames[self.currentFileIndex],mode='r',newline='\n')
            self.csv_reader=csv.reader(self.fp, delimiter=',', quotechar='"')

            if self.skipHeader is True:
                next(self.csv_reader)


readFileNames=[]
readFileNames.append("D:/data/data_01.csv")
readFileNames.append("D:/data/data_02.csv")
readFileNames.append("D:/data/data_03.csv")

my_iterator = MyIterator(readFileNames,skipHeader=True)

for value in my_iterator:
    print(value)

my_iterator = MyIterator(readFileNames,skipHeader=False)

for value in my_iterator:
    print(value)