300 lines
9.3 KiB
Python
300 lines
9.3 KiB
Python
"""
|
||
环境变量加载模块
|
||
|
||
提供统一的环境变量读取接口，支持：
|
||
1. 从 .env 文件加载环境变量
|
||
2. 向后兼容：当 .env 不存在时给出友好提示
|
||
3. 类型转换和默认值处理
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Any, Optional, Union
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 项目根目录
|
||
ROOT_DIR = Path(__file__).resolve().parent
|
||
ENV_FILE_PATH = ROOT_DIR / ".env"
|
||
ENV_EXAMPLE_PATH = ROOT_DIR / ".env.example"
|
||
|
||
|
||
def load_dotenv(env_path: Optional[Path] = None) -> bool:
|
||
"""
|
||
从 .env 文件加载环境变量到 os.environ
|
||
|
||
Args:
|
||
env_path: .env 文件路径，默认为项目根目录的 .env
|
||
|
||
Returns:
|
||
bool: 是否成功加载（文件存在且至少加载了一个变量）
|
||
|
||
注意：
|
||
- 已存在的环境变量不会被覆盖
|
||
- 支持 # 注释和空行
|
||
- 支持格式：KEY=value 或 KEY="value" 或 KEY='value'
|
||
"""
|
||
if env_path is None:
|
||
env_path = ENV_FILE_PATH
|
||
|
||
if not env_path.exists():
|
||
logger.debug(f".env file not found at {env_path} (optional)")
|
||
return False
|
||
|
||
loaded_count = 0
|
||
try:
|
||
with env_path.open("r", encoding="utf-8") as f:
|
||
for line_num, line in enumerate(f, start=1):
|
||
# 去除首尾空白
|
||
line = line.strip()
|
||
|
||
# 跳过空行和注释
|
||
if not line or line.startswith("#"):
|
||
continue
|
||
|
||
# 解析 KEY=VALUE 格式
|
||
if "=" not in line:
|
||
logger.warning(f"Invalid line in {env_path}:{line_num}: {line}")
|
||
continue
|
||
|
||
key, _, value = line.partition("=")
|
||
key = key.strip()
|
||
value = value.strip()
|
||
|
||
# 去除引号（支持单引号和双引号）
|
||
if value and value[0] in ('"', "'") and value[0] == value[-1]:
|
||
value = value[1:-1]
|
||
|
||
# 只设置尚未存在的环境变量（已存在的优先级更高）
|
||
if key and key not in os.environ:
|
||
os.environ[key] = value
|
||
loaded_count += 1
|
||
|
||
logger.info(f"Loaded {loaded_count} environment variables from {env_path}")
|
||
return loaded_count > 0
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to load .env file from {env_path}: {e}")
|
||
return False
|
||
|
||
|
||
def get_env(
|
||
key: str,
|
||
default: Optional[str] = None,
|
||
*,
|
||
required: bool = False,
|
||
hint: Optional[str] = None,
|
||
) -> Optional[str]:
|
||
"""
|
||
读取环境变量（字符串类型）
|
||
|
||
Args:
|
||
key: 环境变量名
|
||
default: 默认值
|
||
required: 是否必需（为 True 且变量不存在时抛出异常）
|
||
hint: 提示信息（当变量不存在时显示）
|
||
|
||
Returns:
|
||
环境变量值或默认值
|
||
|
||
Raises:
|
||
ValueError: 当 required=True 且环境变量不存在时
|
||
|
||
Example:
|
||
>>> get_env("GEMINI_API_KEY", required=True, hint="See .env.example")
|
||
>>> get_env("LOG_LEVEL", default="INFO")
|
||
"""
|
||
value = os.environ.get(key)
|
||
|
||
if value is None:
|
||
if required:
|
||
error_msg = f"Required environment variable '{key}' is not set."
|
||
if hint:
|
||
error_msg += f"\nHint: {hint}"
|
||
elif ENV_EXAMPLE_PATH.exists():
|
||
error_msg += f"\nHint: Copy .env.example to .env and configure it."
|
||
raise ValueError(error_msg)
|
||
return default
|
||
|
||
return value
|
||
|
||
|
||
def get_env_bool(
|
||
key: str,
|
||
default: bool = False,
|
||
*,
|
||
required: bool = False,
|
||
) -> bool:
|
||
"""
|
||
读取环境变量（布尔类型）
|
||
|
||
Args:
|
||
key: 环境变量名
|
||
default: 默认值
|
||
required: 是否必需
|
||
|
||
Returns:
|
||
布尔值
|
||
|
||
认为是 True 的值：true, 1, yes, on（不区分大小写）
|
||
认为是 False 的值：false, 0, no, off（不区分大小写）
|
||
其他值使用默认值
|
||
|
||
Example:
|
||
>>> get_env_bool("DEBUG", default=False)
|
||
>>> get_env_bool("ENABLE_CACHE", default=True)
|
||
"""
|
||
value = get_env(key, required=required)
|
||
|
||
if value is None:
|
||
return default
|
||
|
||
value_lower = value.lower().strip()
|
||
|
||
if value_lower in ("true", "1", "yes", "on"):
|
||
return True
|
||
elif value_lower in ("false", "0", "no", "off"):
|
||
return False
|
||
else:
|
||
logger.warning(
|
||
f"Invalid boolean value for {key}='{value}', using default={default}"
|
||
)
|
||
return default
|
||
|
||
|
||
def get_env_int(
|
||
key: str,
|
||
default: Optional[int] = None,
|
||
*,
|
||
required: bool = False,
|
||
) -> Optional[int]:
|
||
"""
|
||
读取环境变量（整数类型）
|
||
|
||
Args:
|
||
key: 环境变量名
|
||
default: 默认值
|
||
required: 是否必需
|
||
|
||
Returns:
|
||
整数值或默认值
|
||
|
||
Example:
|
||
>>> get_env_int("TIMEOUT", default=30)
|
||
>>> get_env_int("MAX_WORKERS", default=4)
|
||
"""
|
||
value = get_env(key, required=required)
|
||
|
||
if value is None:
|
||
return default
|
||
|
||
try:
|
||
return int(value)
|
||
except ValueError:
|
||
logger.warning(
|
||
f"Invalid integer value for {key}='{value}', using default={default}"
|
||
)
|
||
return default
|
||
|
||
|
||
def get_env_float(
|
||
key: str,
|
||
default: Optional[float] = None,
|
||
*,
|
||
required: bool = False,
|
||
) -> Optional[float]:
|
||
"""
|
||
读取环境变量（浮点数类型）
|
||
|
||
Args:
|
||
key: 环境变量名
|
||
default: 默认值
|
||
required: 是否必需
|
||
|
||
Returns:
|
||
浮点数值或默认值
|
||
|
||
Example:
|
||
>>> get_env_float("THRESHOLD", default=0.5)
|
||
"""
|
||
value = get_env(key, required=required)
|
||
|
||
if value is None:
|
||
return default
|
||
|
||
try:
|
||
return float(value)
|
||
except ValueError:
|
||
logger.warning(
|
||
f"Invalid float value for {key}='{value}', using default={default}"
|
||
)
|
||
return default
|
||
|
||
|
||
def check_required_env_vars(required_keys: list[str]) -> list[str]:
|
||
"""
|
||
检查多个必需的环境变量是否已设置
|
||
|
||
Args:
|
||
required_keys: 必需的环境变量名列表
|
||
|
||
Returns:
|
||
未设置的环境变量名列表（空列表表示全部已设置）
|
||
|
||
Example:
|
||
>>> missing = check_required_env_vars(["GEMINI_API_KEY", "AZURE_API_KEY"])
|
||
>>> if missing:
|
||
>>> print(f"Missing required env vars: {', '.join(missing)}")
|
||
"""
|
||
missing = []
|
||
for key in required_keys:
|
||
if key not in os.environ:
|
||
missing.append(key)
|
||
return missing
|
||
|
||
|
||
def print_env_setup_guide():
|
||
"""
|
||
打印环境变量配置指南（当检测到缺少必需配置时使用）
|
||
"""
|
||
guide = """
|
||
╔════════════════════════════════════════════════════════════════════════════╗
|
||
║ 环境变量配置指南 ║
|
||
╠════════════════════════════════════════════════════════════════════════════╣
|
||
║ 检测到缺少必需的环境变量配置。请按以下步骤配置： ║
|
||
║ ║
|
||
║ 1. 复制模板文件 ║
|
||
║ cp .env.example .env ║
|
||
║ ║
|
||
║ 2. 编辑 .env 文件，填入实际的 API 密钥和配置 ║
|
||
║ vim .env # 或使用其他编辑器 ║
|
||
║ ║
|
||
║ 3. 确保 .env 文件不会提交到版本控制 ║
|
||
║ （已在 .gitignore 中排除） ║
|
||
║ ║
|
||
║ 常用配置项： ║
|
||
║ - GEMINI_API_KEY: Google Gemini API 密钥 ║
|
||
║ - AZURE_API_KEY: Azure OpenAI API 密钥 ║
|
||
║ - QWEN_API_KEY: 阿里云通义千问 API 密钥 ║
|
||
║ - GLM_API_KEY: 智谱 ChatGLM API 密钥 ║
|
||
║ ║
|
||
║ 详细说明请查看 .env.example 文件中的注释 ║
|
||
╚════════════════════════════════════════════════════════════════════════════╝
|
||
"""
|
||
print(guide)
|
||
|
||
|
||
# 自动加载 .env 文件（当模块被导入时）
|
||
_dotenv_loaded = load_dotenv()
|
||
|
||
if not _dotenv_loaded and not ENV_FILE_PATH.exists():
|
||
logger.info(
|
||
f".env file not found at {ENV_FILE_PATH}. "
|
||
f"This is optional but recommended for managing sensitive configuration. "
|
||
f"See .env.example for reference."
|
||
)
|