30. python sys.stderr重定到sys.stdout 并同时将输出归档一份到文件

使用场景

对于远程执行python的时候，调用者可能需要将错误和标准输出都放一个文件
以下是两种实现方式

第一种方式

借助StringIO ， 缺点是任务执行结束才会看到输出，__enter__中将StringIO的数据进行了输出，而__enter__是程序运行结束才会执行，因此日志输出有延迟，优点是代码简单易懂

import json, argparse,sys,traceback
from io import StringIO
def get_error():
    """
    :description: 获取捕获的错误信息
    :return:
    """
    etype, value, tb = sys.exc_info()
    return " ".join(traceback.format_exception(etype, value, tb))


class STDRedirect:
    """
    :description: 输出重定向
    """
    def __init__(self, filename, mode='a+', buffering=100):
        self.o = sys.stdout  # 保存原始的stdout
        self.e = sys.stderr  # 保存原始的stdout
        self.log = open('/root/alg/run.log', mode=mode, buffering=buffering,encoding='utf-8')

    def __enter__(self):
        """
        :description: 重定向stdout到StringIO
        """
        self.cache = StringIO()
        sys.stdout, sys.stderr = self.cache, self.cache
        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, sys.stderr = self.o, self.e
        value = self.cache.getvalue()
        self.log.write(value)
        self.log.flush()
        self.log.close()
        self.cache.close()
        print(value, file=sys.stdout)

"""
    :description: 运行代码
    :param code:
    :return:
"""
if __name__ == "__main__":
# 被with包裹的所有代码，其stdout和stderr都被重指向了StringIO
# 而在__exit__中我们恢复了原始的stdout和stderr,这意味着
# 被with包裹的代码，所有的print都会录入到文件一份
# 不被with包裹的代码则和以前一样
    with STDRedirect("test.log"):
        try:
            """这是你的代码逻辑"""
        except Exception as e:
            print(get_error()

第二种方式

优点，可以实时输出日志，缺点是代码比较复杂。

# -*- coding: UTF-8 -*-
import io
import os
import time
import sys
import threading
import traceback
import ast
import json
import argparse
from datetime import datetime, timedelta

PIPE_BUFFER_SIZE = 4096
# 日志结束标志，当后端收到此标志，意味着程序执行结束
END_FLAG = "[结束]"


def date_string_now(fmt="%Y-%m-%d %H:%M:%S", day=0):
    # 获取当前日期和时间
    time = datetime.now() + timedelta(days=day)

    # 只获取当前日期部分
    current_date = time.strftime(fmt)

    return current_date


def get_error():
    """
    :description: 获取捕获的错误信息
    :return:
    """
    type, value, tb = sys.exc_info()
    return " ".join(traceback.format_exception(type, value, tb))


class STDRedirect:
    def __init__(self, path=None, mode='a+', buffering=100):
        self._err = sys.stderr
        self.log = open(path, mode=mode, buffering=buffering, encoding='utf-8')
        self._exc = None

    def __enter__(self):
        sys.stderr = sys.stdout
        self._pr, self._pw = os.pipe()
        fno = self._original_std = getattr(sys, 'stdout').fileno()
        self._original_std_copy = os.dup(fno)
        os.dup2(self._pw, fno)
        self.watch_fd_thread = threading.Thread(target=self._watch_pipe_fd, daemon=True)
        self.watch_fd_thread.start()
        return self

    def _watch_pipe_fd(self):

        try:
            bts = os.read(self._pr, PIPE_BUFFER_SIZE)
            while bts:
                time.sleep(0.01)
                bts_to_str = bts.decode(errors="replace", encoding="utf-8")
                if "ERROR" not in bts_to_str:
                    string = "\033[32m%s - [RUN_CODE] - INFO - %s\033[0m" % (date_string_now(), bts_to_str)
                else:
                    string = "\033[31m%s - [RUN_CODE] - %s\033[0m" % (date_string_now(), bts_to_str)
                self.write(string)
                os.write(self._original_std_copy, string.encode("utf-8"))
                if END_FLAG in string:
                    break
                bts = os.read(self._pr, PIPE_BUFFER_SIZE)
        except Exception:
            self._exc = sys.exc_info()

    def write(self, string):
        try:
            self.log.write(string.replace('\u0000', ""))
            self.log.flush()
        except Exception:
            os.write(self._original_std_copy, get_error().encode("utf-8"))

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

        if self.watch_fd_thread is not None:
            self.watch_fd_thread.join()
        # restore original FDs
        os.dup2(self._original_std_copy, self._original_std)
        os.close(self._original_std_copy)
        sys.stderr = self._err

        if self._exc:
            etype, value, tb = self._exc
            traceback.print_exception(etype, value, tb)
        if self.log:
            self.log.flush()
            self.log.close()



def create_dir(path):
    directory = os.path.dirname(path)

    if not os.path.exists(directory):

        os.makedirs(directory)



with STDRedirect("a.log"):
    print("哈哈哈")
    print("999")
    print("[结束]")
print("原始的stdout已经恢复")