[with構文の理解・応用] Pythonの標準出力の出力先を柔軟に切り替える

Pythonの標準出力(printやsys.stdout)の出力先（file出力, console出力）をwith構文で適宜切り替えながら使う方法を紹介します。

動機

以下の様にすると標準出力の出力先を処理に応じて変更できますが、

import sys
# 一時的にfile出力に変更
sys.stdout = open('out.log', 'a+')

...
sys.stdout.write('fugahoge')
...

# console出力に戻す
sys.stdout = sys.__stdout__

処理が終わった際に、元に戻す処理が手間なので何かないかなーと考えていて、with構文を使うとfile open()後のclose()が不必要になったことを思い出して、本記事を書くに至りました。

本記事に書いてあること

• 標準出力の変更法
• with構文の説明
• with構文での標準出力の出力先の変更法

with構文

with構文は、以下の様なファイルの読み書きやtensorflowのgradient_tapeなど色々な場面で使われています。

with open('', 'r') as f:
  f.read()

参考: with構文(Python)

with構文では何が起きているか

with構文を使うと何が起きているかというと、withの後にinstanseを指定（もしくは生成）すると、以下のような流れで処理が走ります。
1. そのinstanseの.__enter__()メソッドが呼ばれる
2. with構文内の処理がはしる
3. そのinstanseの.__exit__()メソッドが呼ばれる

コード例

以下のようなコードを実行すると、

class Logger():
    def setIO(self, *args, **kwargs):
        # TestIO instanseの生成
        return TestIO(*args, **kwargs)

class TestIO():
    def __enter__(self):
        print('enter')

    def __exit__(self, *args):
        print('exit')

logger = Logger()
with logger.setIO():
    print('---- in with syntax ----')

以下のように出力されます。

console

enter
---- in with syntax ----
exit

これで、.__enter__()メソッド -> with構文内の処理 -> .__exit__()メソッド、の順に処理がはしっていることが確認できました。

参考:
- with構文とは何なのか - 年中アイス
- [Python] with構文で使用できるクラスを実装する

標準出力の出力先をwith構文で柔軟に切り替える

最後に、本記事の主題である「標準出力の出力先をwith構文で柔軟に切り替える」のコード例を紹介します。

import sys

class SetIO():
    """with構文でI/Oを切り替えるためのクラス"""
    def __init__(self, filename: str):
        self.filename = filename

    def __enter__(self):
        sys.stdout = _STDLogger(out_file=self.filename)

    def __exit__(self, *args):
        sys.stdout = sys.__stdout__

class _STDLogger():
    """カスタムI/O"""
    def __init__(self, out_file='out.log'):
        self.log = open(out_file, "a+")

    def write(self, message):
        self.log.write(message)

    def flush(self):
        # this flush method is needed for python 3 compatibility.
        pass

print('before with block')

with SetIO('out.log'):
    # file出力に切り替え
    print('---- in with syntax ----')

print('after with block')

上記のコードを実行すると、
コンソールには、以下が出力されます。

console

before with block
after with block

そして、ファイル (out.log)には以下が出力されます。

out.log

---- in with syntax ----

上記の方法で標準出力先をwith構文を切り替えることができました。

まとめ

出力先の変更は、組み込みのlogger moduleの名前空間を使えば可能ですが、細かく出力先をかえるには不便だと思っています。（loggerに精通していないだけの可能性はあります）
そこで、本記事で紹介した手法が使える場面もあると考えています。
何かの参考になれば幸いです！

Refs

• with構文とは何なのか - 年中アイス
• [Python] with構文で使用できるクラスを実装する
• with構文(Python)