Spaces:
Running
Running
""" | |
Common utilities. | |
""" | |
from asyncio import AbstractEventLoop | |
from io import BytesIO | |
import base64 | |
import json | |
import logging | |
import logging.handlers | |
import os | |
import platform | |
import sys | |
import time | |
from typing import AsyncGenerator, Generator | |
import warnings | |
import requests | |
LOGDIR = "./" | |
handler = None | |
visited_loggers = set() | |
def build_logger(logger_name, logger_filename): | |
global handler | |
formatter = logging.Formatter( | |
fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s", | |
datefmt="%Y-%m-%d %H:%M:%S", | |
) | |
# Set the format of root handlers | |
if not logging.getLogger().handlers: | |
if sys.version_info[1] >= 9: | |
# This is for windows | |
logging.basicConfig(level=logging.INFO, encoding="utf-8") | |
else: | |
if platform.system() == "Windows": | |
warnings.warn( | |
"If you are running on Windows, " | |
"we recommend you use Python >= 3.9 for UTF-8 encoding." | |
) | |
logging.basicConfig(level=logging.INFO) | |
logging.getLogger().handlers[0].setFormatter(formatter) | |
# Redirect stdout and stderr to loggers | |
stdout_logger = logging.getLogger("stdout") | |
stdout_logger.setLevel(logging.INFO) | |
sl = StreamToLogger(stdout_logger, logging.INFO) | |
sys.stdout = sl | |
stderr_logger = logging.getLogger("stderr") | |
stderr_logger.setLevel(logging.ERROR) | |
sl = StreamToLogger(stderr_logger, logging.ERROR) | |
sys.stderr = sl | |
# Get logger | |
logger = logging.getLogger(logger_name) | |
logger.setLevel(logging.INFO) | |
# Avoid httpx flooding POST logs | |
logging.getLogger("httpx").setLevel(logging.WARNING) | |
# if LOGDIR is empty, then don't try output log to local file | |
if LOGDIR != "": | |
os.makedirs(LOGDIR, exist_ok=True) | |
filename = os.path.join(LOGDIR, logger_filename) | |
handler = logging.handlers.TimedRotatingFileHandler( | |
filename, when="D", utc=True, encoding="utf-8" | |
) | |
handler.setFormatter(formatter) | |
for l in [stdout_logger, stderr_logger, logger]: | |
if l in visited_loggers: | |
continue | |
visited_loggers.add(l) | |
l.addHandler(handler) | |
return logger | |
class StreamToLogger(object): | |
""" | |
Fake file-like stream object that redirects writes to a logger instance. | |
""" | |
def __init__(self, logger, log_level=logging.INFO): | |
self.terminal = sys.stdout | |
self.logger = logger | |
self.log_level = log_level | |
self.linebuf = "" | |
def __getattr__(self, attr): | |
return getattr(self.terminal, attr) | |
def write(self, buf): | |
temp_linebuf = self.linebuf + buf | |
self.linebuf = "" | |
for line in temp_linebuf.splitlines(True): | |
# From the io.TextIOWrapper docs: | |
# On output, if newline is None, any '\n' characters written | |
# are translated to the system default line separator. | |
# By default sys.stdout.write() expects '\n' newlines and then | |
# translates them so this is still cross platform. | |
if line[-1] == "\n": | |
encoded_message = line.encode("utf-8", "ignore").decode("utf-8") | |
self.logger.log(self.log_level, encoded_message.rstrip()) | |
else: | |
self.linebuf += line | |
def flush(self): | |
if self.linebuf != "": | |
encoded_message = self.linebuf.encode("utf-8", "ignore").decode("utf-8") | |
self.logger.log(self.log_level, encoded_message.rstrip()) | |
self.linebuf = "" |