with 文と @contextlib.contextmanager が便利

# はじめに

この記事は [Python Advent Calendar その 2](https://qiita.com/advent-calendar/2017/python_python) の 7 日目の記事です。

僕はふだんは Ruby や JavaScript のコードを書いています。一方 Python は趣味で使っていて、Python 歴は 1 年ほどです。この言語を勉強していて [with](http://docs.python.jp/3/reference/compound_stmts.html#with) 文が気に入ったので、今回はこの with 文とコンテキストマネージャについて語ります。

# with 文について

with 文が Python に導入された経緯については [PEP 343](https://www.python.org/dev/peps/pep-0343) の [Abstract](https://www.python.org/dev/peps/pep-0343/#abstract) で次のように書かれています。

> This PEP adds a new statement "with" to the Python language to make it possible to factor out standard uses of try/finally statements.

`try/finally` の標準的な使い方を抽出する目的で追加したとのことです。

例えば [open()](https://docs.python.jp/3/library/functions.html#open) でファイルを開く場合、最終的にファイルが必ず閉じられることを保証するために `try/finally` を使うことができます。

```py
def divide_by_zero(n):
  return n / 0
```

```py
>>> f = None
>>> try:
...   f = open('hidamari.txt', 'r')
...   divide_by_zero(len(f.read()))
... finally:
...   if f:
...     f.close()
...
ZeroDivisionError: division by zero
>>> print(f.closed)
True
```

ここで with 文を使うと、同じことをよりシンプルなコードで実現できます。

```py
>>> with open('hidamari.txt', 'r') as f:
...   divide_by_zero(len(f.read()))
...
ZeroDivisionError: division by zero
>>> print(f.closed)
True
```

ちなみに、この `with open()` を自作すると次のようなコードになります。

```py
class Reading:
  def __init__(self, filename):
    self.filename = filename

  def __enter__(self):
    self.file = open(self.filename, 'r')

    return self.file

  def __exit__(self, type, value, traceback):
    self.file.close()
```

```py
>>> with Reading('hidamari.txt') as f:
...   divide_by_zero(len(f.read()))
...
ZeroDivisionError: division by zero
>>> print(f.closed)
True
```

with 文に渡されるオブジェクトをコンテキストマネージャと呼びます。上記の with 文の `f` にはコンテキストマネージャの `__enter__` の返り値がバインドされます。そして with ブロックを抜けたときに `__exit__` が呼び出されるという仕組みです。`__enter__` が前処理、`__exit__` が後処理のイメージですね。

毎回 `__enter__` と `__exit__` を持つコンテキストマネージャを定義するのは面倒です。そこで [@contextlib.contextmanager](https://docs.python.jp/3/library/contextlib.html#contextlib.contextmanager) というデコレータを使うと、コンテキストマネージャを簡単に定義できます。

```py
import contextlib


@contextlib.contextmanager
def reading(filename):
  file = open(filename, 'r')
  try:
    yield file
  finally:
    file.close()
```

```py
>>> with reading('hidamari.txt') as f:
...   print(f.read())
...
ひだまり荘で、待ってます。
>>> print(f.closed)
True
```

なお with ブロック内で発生した例外は `yield` を実行している箇所で再送出されます。with ブロック内でユーザがどんなコードを実行するのかはコンテキストマネージャからは予想できません。そのため `yield` は `try/finally` で囲っておく必要があります。

```py
import contextlib


@contextlib.contextmanager
def reading_incompletely(filename):
  file = open(filename, 'r')
  yield file
  file.close()  # 呼ばれない可能性がある。
```

```pycon
>>> with reading_incompletely('hidamari.txt') as f:
...   divide_by_zero(len(f.read()))
...
ZeroDivisionError: division by zero
>>> print(f.closed)  # file が close されていない。
False
```

# もうひとつの例

コンテキストマネージャに前処理と後処理を任せることで、with 文を使うユーザはコードをよりシンプルに書くことができました。そして、コンテキストマネージャにはもうひとつの利点があると僕は考えます。それは、限られた範囲でのみ特別な処理を行う場合です。例えば、既存の組み込み関数にモンキーパッチを当てる場合などです。

次の例では `sys.stdout.write()` を独自の関数に上書きすることで、`print()` を呼び出すたびに出力を上書きできるようにしています。

```countdown.py
import contextlib


@contextlib.contextmanager
def overwriting():
  import sys
  original_write = sys.stdout.write  # 元の関数を保持しておく。

  def overwrite(text):
    original_write('\033[K')  # カーソル位置から行末までを消す。
    original_write('\r')  # カーソルを行頭に移動する。
    original_write(text.rstrip('\n'))  # print() で出力する文字列の末尾の改行を取り除く。

  sys.stdout.write = overwrite  # 関数を入れ替える。

  try:
    yield
  finally:
    sys.stdout.write = original_write  # 関数を元に戻す。忘れると大変。
    print('※ sys.stdout.write を元に戻しました。')


if __name__ == '__main__':
  import time

  # この with ブロックだけ print() の挙動が特殊になる。
  with overwriting():
    print('カウントダウン開始！')
    time.sleep(0.4)
    for i in range(10, -1, -1):
      print(i)
      time.sleep(0.4)

  print('')
```

![anime.gif](https://qiita-image-store.s3.amazonaws.com/0/6697/6ee40289-aa95-c3ba-57ff-b100f452496d.gif)

このように限られた範囲でのみ特別な処理を行う場合に、`@contextlib.contextmanager` を使って独自のコンテキストマネージャを作成し、with 文を使うことで

* with 文のシンタックスによって、限定されている範囲が分かりやすい。
* 重要な後処理を必ず実行することを保証できる。

というメリットがあると思います。

# 参考

* [Python ドキュメント » 29.6. contextlib — with 文コンテキスト用ユーティリティ](https://docs.python.jp/3/library/contextlib.html)
* [PEP 343 -- The "with" Statement](https://www.python.org/dev/peps/pep-0343)