# log_utils.py
import logging
from io import StringIO
from flask import g


def init_request_logging(logger):
    # 创建内存缓冲区
    g.log_stream = StringIO()

    # 创建自定义日志处理器
    handler = logging.StreamHandler(g.log_stream)
    handler.setFormatter(logging.Formatter(
        '%(asctime)s [%(levelname)s] %(message)s'
    ))

    # 临时添加到logger
    g.log_handler = handler
    logger.addHandler(handler)


def teardown_request_logging(response, logger):
    if hasattr(g, 'log_handler'):
        logger.removeHandler(g.log_handler)
    return response