pythonで簡易先読みiterator作成

2023年1月20日2023年1月21日

表題のとおりの小ネタです。

テキストファイルの解析をしていると、次の塊の先頭が前の塊の終端を表すことがあり、よく先読みをしたくなったりします。そんなときに「1つだけ」先を読めるiteratorを書いてみました。

Contents

1. 背景
2. 目的
3. iteratorとは…
4. 実装
5. まとめ

背景

なんとなくパッチを扱う処理を書いてたときに欲しくなったので、書いてみました。

目的

pythonで1つだけ先読み(peek)できるiteratorの作成

iteratorとは…

コンテナなど同じ形のデータを繰り返し走査するときに使うものです。pythonでは例えばfor ～ in 〇〇の〇〇の部分に記述することができます。もしlistのようなiterableなオブジェクトが指定された場合は、自動的にiteratorが生成されて返されています。

>>> for i in [1,2,3]:
...     print(i)
... 
1
2
3
>>> for i in iter([1,2,3]):
...     print(i)
... 
1
2
3
>>>

pythonのiteratorはこんな形に定義されています。

https://docs.python.org/3.10/library/stdtypes.html#typeiter

この定義では__iter__()と__next__()を持っていればiteratorということになります。

built-in関数のnext()の引数に指定することで__next__()が呼び出され、繰り返し走査が可能です。終端ではStopIteration 例外をあげる仕様です。built-in関数のiter()の引数にiterableを指定すると、新しいiteratorが返ります。引数にiteratorを指定するとそのiterator自身が返ります。

>>> a = [1,2,3]
>>> i = iter(a)
>>> j = iter(a)
>>> next(i)
1
>>> next(j)
1
>>> i = iter(a)
>>> j = iter(i)
>>> next(i)
1
>>> next(j)
2
>>>

今回はこの単純な iterator に1個だけ先読み(peek）を付けようという話です。

実装

class PeekableIterator:
    def __init__(self, it):
        self.iter = it
        self.next_element = None
        self.got_next = False
    def __iter__(self):
        return self
    def __next__(self):
        self.peek()
        self.got_next = False
        return self.next_element
    def __bool__(self):
        try:
            self.peek()
            return True
        except StopIteration:
            return False
    def peek(self):
        if not self.got_next:
            self.next_element = next(self.iter)
            self.got_next = True
        return self.next_element

if __name__ == '__main__':
    with open('hoge.py', 'rt') as f:
        it = PeekableIterator(f)
        for line in it:
            print(line, end='')
            if it:
                print('next is probably...', it.peek(), end='')

1行読み込んだ後、次の行の有無を見て、あればその行を先にprintしています。ただしpeek()しているだけなので次に進んでいません。次のforループを回さずに先に一行だけ読んでいるということです。

※__bool__()はbool型へのキャスト操作です。

https://docs.python.org/3/reference/datamodel.html#object.bool