prompt_colors

This commit is contained in:
Divided by Zer0
2022-09-11 23:08:17 +02:00
parent ce2d1ff654
commit e29f6c94d3

View File

@ -2,22 +2,32 @@ import sys
from functools import partialmethod from functools import partialmethod
from loguru import logger from loguru import logger
STDOUT_LEVELS = ["GENERATION", "PROMPT"]
def is_stdout_log(record): def is_stdout_log(record):
if record["level"].name == "GENERATION": if record["level"].name in STDOUT_LEVELS:
return(True) return(True)
return(False) return(False)
def is_stderr_log(record): def is_stderr_log(record):
if record["level"].name != "GENERATION": if record["level"].name not in STDOUT_LEVELS:
return(True) return(True)
return(False) return(False)
logfmt = "<level>{level: <10}</level> | <green>{name}</green>:<green>{function}</green>:<green>{line}</green> - <level>{message}</level>" logfmt = "<level>{level: <10}</level> | <green>{name}</green>:<green>{function}</green>:<green>{line}</green> - <level>{message}</level>"
genfmt = "<level>{level: <10}</level> @ <green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{message}</level>" genfmt = "<level>{level: <10}</level> @ <green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{message}</level>"
new_level = logger.level("GENERATION", no=21, color="<cyan>") promptfmt = "<level>{level: <10}</level> @ <green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{message}</level>"
new_level = logger.level("GENERATION", no=22, color="<cyan>")
new_level = logger.level("PROMPT", no=21, color="<yellow>")
logger.__class__.generation = partialmethod(logger.__class__.log, "GENERATION") logger.__class__.generation = partialmethod(logger.__class__.log, "GENERATION")
logger.__class__.prompt = partialmethod(logger.__class__.log, "PROMPT")
config = { config = {
"handlers": [ "handlers": [
{"sink": sys.stderr, "format": logfmt, "colorize":True, "filter": is_stderr_log}, {"sink": sys.stderr, "format": logfmt, "colorize":True, "filter": is_stderr_log},
{"sink": sys.stdout, "format": genfmt, "level": "GENERATION", "colorize":True, "filter": is_stdout_log}, {"sink": sys.stdout, "format": genfmt, "level": "PROMPT", "colorize":True, "filter": is_stdout_log},
], ],
} }
logger.configure(**config) logger.configure(**config)