ざっくりPythonのおさらいその4。contextlib系 #Python

前回に引き続いて、contextlib周りのおさらいです。
その他のおさらい事項にについては、以下ご参照ください。

[ざっくりPythonのおさらいその1。基本文法など]
https://qiita.com/48hands/items/31b7ac2b49addbb8658c

[ざっくりPythonのおさらいその2。おもにデータ保存]
https://qiita.com/48hands/items/ab86b7463268e669d216

[ざっくりPythonのおさらいその3。Numpy, Pandas]
https://qiita.com/48hands/items/b1a71000533e129b27f2

contextlib.contextmanager

contextlib.contextmanagerはデコレータに関連したライブラリ。
withステートメントを使って、簡単にデコレーションできるようになる。

まずはじめに、デコレータの復習。
以下がデコレーションの基本。デコレーション対象の関数fにデコレーションを施す関数を@tagで指定している。

import contextlib


def tag(f):
    def _wrapper(content):
        print("decorated start")
        r = f(content)
        print("decorated end")
        return r

    return _wrapper


@tag
def f(content):
    print(content)


if __name__ == '__main__':
    f('this is test message')

実行結果は以下。

decorated start
this is test message
decorated end

ここからデコレーションを施す関数tagに引数を渡したい場合には、以下のように記述する。

import contextlib


def tag(name):
    def _tag(f):
        def _wrapper(content):
            print("{} start".format(name))
            r = f(content)
            print('{} end'.format(name))
            return r

        return _wrapper

    return _tag


@tag('decorated')
def f(content):
    print(content)


if __name__ == '__main__':
    f('this is test message')

これをわかりやすくするために、contextlib.contextmanagerを使う。contextlib.contextmanagerで書くと以下のようになる。

yieldを使ってデコレーション対象の関数fを呼び出している。

@contextlib.contextmanager
def tag(name):
    print("{} start".format(name))
    yield
    print('{} end'.format(name))


@tag('decorated')
def f(content):
    print(content)


if __name__ == '__main__':
    f('this is test message')

さらに...
ここで、関数fを使っているが、fを使わずに、以下のように記述できる。

import contextlib


@contextlib.contextmanager
def tag(name):
    print("{} start".format(name))
    yield
    print('{} end'.format(name))


if __name__ == '__main__':
    with tag('decorated'):
        print('this is test message')

contextlib.managerはデコレータとしての役割をwithステートメントを使って果たすようなことができている。
とてもわかりやすい！

さらに、withステートメントをネストするようなことも簡単にできる。

import contextlib


@contextlib.contextmanager
def tag(name):
    print("{} start".format(name))
    yield
    print('{} end'.format(name))


if __name__ == '__main__':
    with tag('decorated'):
        print('this is test message')
        with tag('decorated2'):
            print('this is nested message')

実行すると以下のようになる。

decorated start
this is test message
decorated2 start
this is nested message
decorated2 end
decorated end

contextlib.ContextDecorator

用途はcontextlib.contextmanagerと似ている気がするが、実装方法が異なる。ContextDecoratorを継承したクラスを定義して利用する。

import contextlib


# contextlib.ContextDecoratorを継承したクラスを定義
class tag(contextlib.ContextDecorator):
    def __init__(self, name):
        self.name = name
        self.start_tag = '<{}>'.format(name)
        self.end_tag = '</{}>'.format(name)

    # デコレーションに入ったときの動作を定義
    def __enter__(self):
        print(self.start_tag)

    # デコレーションを抜けるときの動作を定義
    def __exit__(self, exc_type, exc_val, exc_tb):
        print(self.end_tag)


if __name__ == '__main__':
    with tag('decorated'):
        print('this is test message')
        with tag('decorated2'):
            print('this is nested message')

<decorated>
this is test message
<decorated2>
this is nested message
</decorated2>
</decorated>

def __exit__(self, exc_type, exc_val, exc_tb):のexec_type, exec_val, exec_tbに何が入っているか確認してみる。

import contextlib


# contextlib.ContextDecoratorを継承したクラスを定義
class tag(contextlib.ContextDecorator):
    def __init__(self, name):
        self.name = name
        self.start_tag = '<{}>'.format(name)
        self.end_tag = '</{}>'.format(name)

    # デコレーションに入ったときの動作を定義
    def __enter__(self):
        print(self.start_tag)

    # デコレーションを抜けるときの動作を定義
    # Exceptionなどの情報が引数に入っている。
    def __exit__(self, exc_type, exc_val, exc_tb):
        print(self.end_tag)
        print('Exception type: {}'.format(exc_type))
        print('Exception value: {}'.format(exc_val))
        print('Exception trace back {}'.format(exc_tb))


if __name__ == '__main__':
    with tag('decorated'):
        raise Exception('Error!!')
        print('hoge hoge')

<decorated>
Traceback (most recent call last):
</decorated>
Exception type: <class 'Exception'>
Exception value: Error!!
Exception trace back <traceback object at 0x109d01f88>
  File "/Users/nagakuray/Desktop/TDD-practice/lesson/context.py", line 26, in <module>
    raise Exception('Error!!')
Exception: Error!!

contextlib.supress

存在しないファイル/path/to/somefile.tmpを例として、
以下のようなコードで例外FileNotFoundErrorを抑圧したいような場合。

import os

try:
    os.remove('/path/to/somefile.tmp')
except FileNotFoundError:
    pass

上記のコードは、contextlib,suppressを使ってwithステートメントで以下のように抑圧(suppress)できる。

import contextlib
import os

with contextlib.suppress(FileNotFoundError):
    os.remove('/path/to/somefile.tmp')

標準入力、標準出力、標準エラー出力、出力のリダイレクト

標準入力

キーボードの入力を待ち受けて出力する。

import sys

for line in sys.stdin:
    print(line)

ループして標準入力を待受て出力する場合はこちら。

import sys

for line in sys.stdin:
    print(line)

標準出力

print('hello')

これを以下のように書ける。

import sys


sys.stdout.write('hello')

標準エラー出力

import sys
import logging

logging.error('Error!')
sys.stderr.write('Error!!')

出力のリダイレクト

contextlib.redirect_stdoutで標準出力をリダイレクトしてファイルに書き出しもできる。Python3.5からの機能。

import contextlib

with open('stdout.log', 'w') as f:
    with contextlib.redirect_stdout(f):
        print('hello')

標準エラー出力をリダイレクトしてファイルに書き出ししたい場合は、contextlib.redirect_stderrを使う。

import contextlib

with open('stderr.log', 'w') as f:
    with contextlib.redirect_stderr(f):
        print('hello')

contextlib.ExitStack

ExitStackを使うと、
callbackが明示的に記述できて、必ず実行したい処理が一目見てわかる。

以下のような処理が合った場合を考える。

def is_ok_job():
    try:
        print('do something')
        raise Exception('Error!')
        return True
    except Exception:
        return False


def cleanup():
    print('clean up')


try:
    is_ok = is_ok_job()
    print('more task')
finally:
    # is_okがTrueでなかった場合にcleanupの処理を実行する
    if not is_ok:
        cleanup()

実行結果は以下。

do something
more task
clean up

このtry~finallyを置き換えるのが、contextlib.ExitStack
以下のように使う。

import contextlib

def is_ok_job():
    try:
        print('do something')
        raise Exception('Error!')
        return True
    except Exception:
        return False


def cleanup():
    print('clean up')


with contextlib.ExitStack() as stack:
    # cleanupをスタックしておく。
    stack.callback(cleanup)

    is_ok = is_ok_job()
    print('more more more')

    # is_okがTrueの場合にpop_allでスタックを空にする。
    # もしFalseの場合はstackされたcleanupがコールバックされる。
    if is_ok:
        stack.pop_all()

実行結果は以下。

do something
more more more
clean up

cleanup関数が引数を持っていて引数を渡す場合は、以下のように書ける。

import contextlib

def is_ok_job():
    try:
        print('do something')
        raise Exception('Error!')
        return True
    except Exception:
        return False


def cleanup(msg, *args, **kwargs):
    print(msg)
    print(args)
    print(kwargs)


with contextlib.ExitStack() as stack:
    # 引数を渡す場合は、以下のように渡す。
    stack.callback(cleanup, 'hello', 'good', 'ng', mysql='3306',
                   postgres='5432')

    is_ok = is_ok_job()
    print('more more more')

    if is_ok:
        stack.pop_all()

実行結果は以下。

do something
more more more
hello
('good', 'ng')
{'mysql': '3306', 'postgres': '5432'}

また、withステートメントの中でインナー関数としてcleanup関数を記述できる。
@stack.callbackをインナー関数に記述する。

with contextlib.ExitStack() as stack:
    @stack.callback
    def inner_cleanup():
        print('inner cleanup')


    is_ok = is_ok_job()
    print('more more more')

    if is_ok:
        stack.pop_all()

実行結果は以下。

do something
more more more
inner cleanup

callbackが明示的に記述されることによって、必ず実行したい処理が一目見てわかるのがExitStackの利点といえる。

ざっくりPythonのおさらい その4。contextlib系

contextlib.contextmanager

contextlib.ContextDecorator

contextlib.supress

標準入力、標準出力、標準エラー出力、出力のリダイレクト

標準入力

標準出力

標準エラー出力

出力のリダイレクト

contextlib.ExitStack

ざっくりPythonのおさらいその4。contextlib系