24. python使用上下文管理器

第一种方式

import io
import os

import sys, threading
from contextlib import contextmanager

PIPE_BUFFER_SIZE = 1000


class OutStream():
    def __init__(self, name, watchfd=True):
        self.name = name
        self.watchfd = watchfd
        self._setup_stream_redirects(name)

    def _watch_pipe_fd(self):
        try:
            bts = os.read(self._fid, PIPE_BUFFER_SIZE)
            while bts:
                self.write(bts.decode(errors="replace"))
                os.write(self._original_stdstream_copy, bts)
                bts = os.read(self._fid, PIPE_BUFFER_SIZE)
        except Exception:
            self._exc = sys.exc_info()
    def write(self, data):
        self.echo.write(data)
        self.echo.flush()
        with open("%s.log" % self.name, "a+") as f:
            f.write(data)
    def _setup_stream_redirects(self, name):
        """
        os.dup2(piper_write, fno) 是将piper_write的副本复制到fno时的，fno重定向到piper_write
        因为底层操作的是副本，因此piper_write之后就可以直接调用os.close(piper_write)关闭
        """
        pip_reader, piper_write = os.pipe()
        fno = self._original_stdstream_fd = getattr(sys, name).fileno()
        self._original_stdstream_copy = os.dup(fno)
        os.dup2(piper_write, fno)  # 原本写入到stdout的数据，此后都会写入到pip_reader
        os.close(piper_write)# 重定向完成后
        self._fid = pip_reader  # stdout重定向到pip_reader
        self.echo = io.TextIOWrapper(io.FileIO(self._original_stdstream_copy,"w"),encoding="utf-8")
        self.watch_fd_thread = threading.Thread(target=self._watch_pipe_fd)
        self.watch_fd_thread.daemon = True
        self.watch_fd_thread.start()
    def close(self):
        os.dup2(self._original_stdstream_copy, self._original_stdstream_fd)# 恢复系统stdout 和stderr
        os.close(self._original_stdstream_copy) # 关闭此文件描述符




ou = OutStream("stdout", watchfd=False)
er= OutStream("stderr", watchfd=False)

@contextmanager
def _redirected_io():
    sys_stdout, sys_stderr = sys.stdout, sys.stderr
    try:
        sys.stdout, sys.stderr = ou, er
        yield
    finally:
        ou.close()
        er.close()
        sys.stdout, sys.stderr = sys_stdout, sys_stderr

with _redirected_io():
    print("hello info")
    print("hello error", file=sys.stderr)

print("结束")

第二种方式

import sys
from io import StringIO

class TeeStdout:
    def __init__(self, filename, mode='w', buffering=100):
        self.terminal = sys.stdout  # 保存原始的stdout
        self.log = open(filename, mode=mode, buffering=buffering)

    def __enter__(self):
        """
        :description: 重定向stdout到StringIO
        """
        self.io = StringIO()
        sys.stdout = self.io
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """

        :param exc_type:
        :param exc_val:
        :param exc_tb:
        1. 恢复系统stdout
        2. 写入文件
        3. 重新打印到恢复后的stdout
        """
        sys.stdout = self.terminal
        value = self.io.getvalue()
        self.log.write(value)  #
        self.log.close()
        print("捕获："+ value)


# 使用上下文管理器
with TeeStdout('output.log') as tee:
    print("This will be written to both the terminal and the log file.")