#!/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()