複数のresourceにループで触る時にcontextlib.ExitStackが便利という話。

複数のresourceにループで触る時にcontextlib.ExitStackが便利という話のメモ。 例えば以下の様な条件で動作する処理を書きたいことがあるとする。

1. 対象となるresource(ファイルなど)を複数取る
2. そのうちN行までを取り出す(それ以降は読まない)
3. 取り出した行だけをprintする

(あるいは対象となる全体でN行だけ取り出すように変えても良いかもしれない)

ng

context managerをloan patternのように使うのはpythonでの標準的な書き方。なのでloop中にwith構文を使ってresourceから読み出すようにする。返されるのはiterableなオブジェクトなのでitertoolsのisliceを使ってN行だけ通りだす。

それなりにまともなコードなような気がする。

ng code

import itertools
import contextlib

def peek(files, *, n):
    r = []
    for f in files:
        with contextlib.closing(f) as rf:
            r.append(itertools.islice(rf, n))
    for line in itertools.chain.from_iterable(r):
        print(line.rstrip())

ぱっと見た形ではこれで良いかのように思うがダメ。

from io import StringIO

# ちょっと手抜きでStringIOを使っているけれど。ファイルを取るのも実質同じ。

lines = [str(x) for x in range(10)]
files = [StringIO("\n".join(lines)) for i in range(5)]
peek(files, n=3)

以下の様なエラーが出る。

ValueError: I/O operation on closed file

理由は、返されるiterable objectを利用する前にresourceを開放しているため。

ok (with contextlib.ExitStack)

このようなときにcontexttlib.ExitStackが便利に使えることがある。以下の様にコードを書き換える。

--- 00ng.py  2019-02-03 22:08:33.971231688 +0900
+++ 01ok.py   2019-02-03 22:09:21.187966956 +0900
@@ -4,12 +4,13 @@
 
 
 def peek(files, *, n):
-    r = []
-    for f in files:
-        with contextlib.closing(f) as rf:
+    with contextlib.ExitStack() as s:
+        r = []
+        for f in files:
+            rf = s.enter_context(contextlib.closing(f))
             r.append(itertools.islice(rf, n))
-    for line in itertools.chain.from_iterable(r):
-        print(line.rstrip())
+        for line in itertools.chain.from_iterable(r):
+            print(line.rstrip())

今度は実行可能になる。

$ python 01ok.py
0
1
2
0
1
2
0
1
2
0
1
2
0
1
2

contextlib.ExitStackは構文上に現れるwithをオブジェクトの操作に変えられるのでこういうことができる(あとはgoのdeferのように使うこともできるけれど。これは別の話。)。