neollm/utils/utils.py

from __future__ import annotations

import os
from typing import Any

from typing_extensions import TypedDict

from neollm.types import PrintColor


class CPrintParam(TypedDict, total=False):
    text: Any
    color: PrintColor | None
    background: bool
    light: bool
    bold: bool
    italic: bool
    underline: bool
    kwargs: dict[str, Any]


def cprint(
    *text: Any,
    color: PrintColor | None = None,
    background: bool = False,
    light: bool = False,
    bold: bool = False,
    italic: bool = False,
    underline: bool = False,
    kwargs: dict[str, Any] = {},
) -> None:
    """
    色付けなどリッチにprint

    Args:
        *text: 表示するテキスト。
        color (PrintColor): テキストの色: 'black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'。
        background (bool): 背景色
        light (bool): 淡い色にするか
        bold (bool): 太字
        italic (bool): 斜体
        underline (bool): 下線
        **kwargs: printの引数
    """
    # ANSIエスケープシーケンスを使用して、テキストを書式設定して表示する
    format_string = ""

    # 色の設定
    color2code: dict[PrintColor, int] = {
        "black": 30,
        "red": 31,
        "green": 32,
        "yellow": 33,
        "blue": 34,
        "magenta": 35,
        "cyan": 36,
        "white": 37,
    }
    if color is not None and color in color2code:
        code = color2code[color]
        if background:
            code += 10
        elif light:
            code += 60
        format_string += f"\033[{code}m"
    if bold:
        format_string += "\033[1m"
    if italic:
        format_string += "\033[3m"
    if underline:
        format_string += "\033[4m"

    # テキストの表示
    for text_i in text:
        print(format_string + str(text_i) + "\033[0m", **kwargs)


def ensure_env_var(var_name: str | None = None, default: str | None = None) -> str:
    if var_name is None:
        return ""
    if os.environ.get(var_name, "") == "":
        if default is None:
            raise ValueError(f"{var_name}をenvで設定しよう")
        cprint(f"WARNING: {var_name}が設定されていません。{default}を使用します。", color="yellow", background=True)
        os.environ[var_name] = default
    return os.environ[var_name]


def suport_unrecomended_env_var(old_key: str, new_key: str) -> None:
    """非推奨の環境変数をサポートする

    Args:
        old_key (str): 非推奨の環境変数名
        new_key (str): 推奨の環境変数名
    """
    if os.getenv(old_key) is not None and os.getenv(new_key) is None:
        cprint(f"WARNING: {old_key}ではなく、{new_key}にしてね", color="yellow", background=True)
        os.environ[new_key] = os.environ[old_key]