playwrite/config/env_loader.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
环境变量加载器

使用 python-dotenv 加载 .env 文件，并提供类型转换功能。
"""

import os
from pathlib import Path
from typing import Any, Optional, Type, TypeVar
from dotenv import load_dotenv

# 项目根目录
PROJECT_ROOT = Path(__file__).parent.parent


def load_env_file(env_file: Optional[str] = None) -> None:
    """
    加载 .env 文件

    Args:
        env_file: .env 文件路径，默认为项目根目录下的 .env
    """
    if env_file is None:
        env_file = PROJECT_ROOT / ".env"
    else:
        env_file = Path(env_file)

    load_dotenv(env_file)


def get_env(key: str, default: Any = None) -> str:
    """
    获取环境变量

    Args:
        key: 环境变量名
        default: 默认值

    Returns:
        环境变量值
    """
    return os.getenv(key, default)


def get_env_bool(key: str, default: bool = False) -> bool:
    """
    获取布尔类型环境变量

    Args:
        key: 环境变量名
        default: 默认值

    Returns:
        布尔值
    """
    value = os.getenv(key, "")
    if not value:
        return default
    return value.lower() in ("true", "1", "yes", "on")


def get_env_int(key: str, default: int = 0) -> int:
    """
    获取整数类型环境变量

    Args:
        key: 环境变量名
        default: 默认值

    Returns:
        整数值
    """
    value = os.getenv(key, "")
    if not value:
        return default
    try:
        return int(value)
    except ValueError:
        return default


def get_env_float(key: str, default: float = 0.0) -> float:
    """
    获取浮点数类型环境变量

    Args:
        key: 环境变量名
        default: 默认值

    Returns:
        浮点数值
    """
    value = os.getenv(key, "")
    if not value:
        return default
    try:
        return float(value)
    except ValueError:
        return default


def set_env(key: str, value: Any) -> None:
    """
    设置环境变量（仅在当前进程中有效）

    Args:
        key: 环境变量名
        value: 环境变量值
    """
    os.environ[key] = str(value)


def save_env_file(
    env_file: Optional[str] = None, env_dict: Optional[dict] = None
) -> bool:
    """
    保存环境变量到 .env 文件

    Args:
        env_file: .env 文件路径，默认为项目根目录下的 .env
        env_dict: 要保存的环境变量字典，如果为 None 则保存当前所有环境变量

    Returns:
        保存是否成功
    """
    if env_file is None:
        env_file = PROJECT_ROOT / ".env"
    else:
        env_file = Path(env_file)

    try:
        # 确保目录存在
        env_file.parent.mkdir(parents=True, exist_ok=True)

        # 读取现有的 .env 文件以保留注释
        existing_lines = []
        if env_file.exists():
            with open(env_file, "r", encoding="utf-8") as f:
                existing_lines = f.readlines()

        # 如果提供了 env_dict，则保存指定的环境变量
        if env_dict is not None:
            # 构建新的文件内容
            new_content = []
            processed_keys = set()

            for line in existing_lines:
                stripped = line.strip()
                # 保留注释和空行
                if not stripped or stripped.startswith("#"):
                    new_content.append(line)
                # 更新已存在的键值对
                elif "=" in stripped and not stripped.startswith("#"):
                    key = stripped.split("=")[0].strip()
                    if key in env_dict:
                        value = env_dict[key]
                        # 处理布尔值的格式
                        if isinstance(value, bool):
                            value = "true" if value else "false"
                        new_content.append(f"{key}={value}\n")
                        processed_keys.add(key)
                    else:
                        new_content.append(line)

            # 添加新的键值对
            for key, value in env_dict.items():
                if key not in processed_keys:
                    # 处理布尔值的格式
                    if isinstance(value, bool):
                        value = "true" if value else "false"
                    new_content.append(f"{key}={value}\n")

            # 写入文件
            with open(env_file, "w", encoding="utf-8") as f:
                f.writelines(new_content)
        else:
            # 如果没有提供 env_dict，则不执行任何操作
            # 因为保存所有环境变量可能会包含系统变量
            return False

        return True
    except IOError as e:
        print(f"保存 .env 文件失败: {e}")
        return False


def update_env_file(env_file: Optional[str] = None, **kwargs) -> bool:
    """
    更新 .env 文件中的特定环境变量

    Args:
        env_file: .env 文件路径
        **kwargs: 要更新的环境变量键值对

    Returns:
        更新是否成功
    """
    return save_env_file(env_file, kwargs)


# 自动加载 .env 文件
load_env_file()