Files
2026-07-13 12:34:57 +08:00

341 lines
7.8 KiB
Python

"""
公共工具函数
包含常用的工具函数和辅助方法
"""
import hashlib
import secrets
import socket
import re
import uuid
from datetime import datetime, timezone
from typing import Optional, Union
from loguru import logger as loguru_logger
logger = loguru_logger.bind(name="utils")
def generate_id() -> str:
"""生成唯一ID"""
return str(uuid.uuid4())
def get_local_ip() -> str:
"""获取本机真实 IP 地址"""
try:
# 连接到一个外部地址来获取本地 IP
# 不实际发送数据,只是用于获取本地 IP
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
# 连接到公共 DNS(不需要实际连接成功)
s.connect(('8.8.8.8', 80))
ip = s.getsockname()[0]
finally:
s.close()
return ip
except Exception:
# 如果失败,尝试使用主机名
try:
hostname = socket.gethostname()
ip = socket.gethostbyname(hostname)
# 如果不是有效的 IP(可能是 127.0.0.1),返回 localhost
if ip.startswith('127.'):
return 'localhost'
return ip
except Exception:
return 'localhost'
def generate_short_id(length: int = 8) -> str:
"""生成短ID"""
return secrets.token_urlsafe(length)
def format_datetime(dt: datetime, format_str: str = "%Y-%m-%d %H:%M:%S") -> str:
"""
格式化日期时间
Args:
dt: 日期时间对象
format_str: 格式字符串
Returns:
str: 格式化后的字符串
"""
if dt is None:
return ""
return dt.strftime(format_str)
def parse_datetime(date_str: str, format_str: str = "%Y-%m-%d %H:%M:%S") -> Optional[datetime]:
"""
解析日期时间字符串
Args:
date_str: 日期时间字符串
format_str: 格式字符串
Returns:
Optional[datetime]: 解析后的日期时间对象
"""
try:
return datetime.strptime(date_str, format_str)
except ValueError:
logger.warning(f"无法解析日期时间: {date_str}")
return None
def validate_email(email: str) -> bool:
"""
验证邮箱格式
Args:
email: 邮箱地址
Returns:
bool: 是否有效
"""
pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
return bool(re.match(pattern, email))
def validate_phone(phone: str) -> bool:
"""
验证手机号格式
Args:
phone: 手机号
Returns:
bool: 是否有效
"""
pattern = r'^1[3-9]\d{9}$'
return bool(re.match(pattern, phone))
def hash_password(password: str) -> str:
"""
哈希密码
Args:
password: 原始密码
Returns:
str: 哈希后的密码
"""
salt = secrets.token_hex(16)
password_hash = hashlib.pbkdf2_hmac(
'sha256',
password.encode('utf-8'),
salt.encode('utf-8'),
100000
)
return f"{salt}:{password_hash.hex()}"
def verify_password(password: str, password_hash: str) -> bool:
"""
验证密码
Args:
password: 原始密码
password_hash: 哈希后的密码
Returns:
bool: 是否匹配
"""
try:
salt, hash_value = password_hash.split(':')
password_hash_check = hashlib.pbkdf2_hmac(
'sha256',
password.encode('utf-8'),
salt.encode('utf-8'),
100000
)
return password_hash_check.hex() == hash_value
except ValueError:
return False
def generate_token(length: int = 32) -> str:
"""
生成随机令牌
Args:
length: 令牌长度
Returns:
str: 随机令牌
"""
return secrets.token_urlsafe(length)
def generate_api_key() -> str:
"""生成API密钥"""
return f"pk_{secrets.token_urlsafe(32)}"
def mask_sensitive_data(data: str, mask_char: str = "*", visible_chars: int = 4) -> str:
"""
遮蔽敏感数据
Args:
data: 原始数据
mask_char: 遮蔽字符
visible_chars: 可见字符数
Returns:
str: 遮蔽后的数据
"""
if len(data) <= visible_chars:
return mask_char * len(data)
return data[:visible_chars] + mask_char * (len(data) - visible_chars)
def safe_get(data: dict, key: str, default: any = None) -> any:
"""
安全获取字典值
Args:
data: 字典数据
key: 键,支持点号分隔的嵌套键
default: 默认值
Returns:
any: 值或默认值
"""
keys = key.split('.')
value = data
try:
for k in keys:
value = value[k]
return value
except (KeyError, TypeError):
return default
def chunk_list(lst: list, chunk_size: int) -> list:
"""
将列表分块
Args:
lst: 原始列表
chunk_size: 块大小
Returns:
list: 分块后的列表
"""
return [lst[i:i + chunk_size] for i in range(0, len(lst), chunk_size)]
def remove_none_values(data: dict) -> dict:
"""
移除字典中的None值
Args:
data: 原始字典
Returns:
dict: 清理后的字典
"""
return {k: v for k, v in data.items() if v is not None}
def deep_merge_dict(dict1: dict, dict2: dict) -> dict:
"""
深度合并字典
Args:
dict1: 第一个字典
dict2: 第二个字典
Returns:
dict: 合并后的字典
"""
result = dict1.copy()
for key, value in dict2.items():
if key in result and isinstance(result[key], dict) and isinstance(value, dict):
result[key] = deep_merge_dict(result[key], value)
else:
result[key] = value
return result
def get_current_timestamp() -> int:
"""获取当前时间戳(毫秒)"""
return int(datetime.now(timezone.utc).timestamp() * 1000)
def format_file_size(size_bytes: int) -> str:
"""
格式化文件大小
Args:
size_bytes: 字节数
Returns:
str: 格式化后的大小
"""
if size_bytes == 0:
return "0 B"
size_names = ["B", "KB", "MB", "GB", "TB"]
i = 0
while size_bytes >= 1024 and i < len(size_names) - 1:
size_bytes /= 1024.0
i += 1
return f"{size_bytes:.1f} {size_names[i]}"
def is_valid_url(url: str) -> bool:
"""
验证URL格式
Args:
url: URL字符串
Returns:
bool: 是否有效
"""
pattern = r'^https?://(?:[-\w.])+(?:\:[0-9]+)?(?:/(?:[\w/_.])*(?:\?(?:[\w&=%.])*)?(?:\#(?:[\w.])*)?)?$'
return bool(re.match(pattern, url))
class RetryHelper:
"""重试辅助类"""
def __init__(self, max_retries: int = 3, delay: float = 1.0, backoff_factor: float = 2.0):
self.max_retries = max_retries
self.delay = delay
self.backoff_factor = backoff_factor
def execute(self, func, *args, **kwargs):
"""执行带重试的函数"""
import time
last_exception = None
current_delay = self.delay
for attempt in range(self.max_retries + 1):
try:
return func(*args, **kwargs)
except Exception as e:
last_exception = e
if attempt < self.max_retries:
logger.warning(f"执行失败,{current_delay}秒后重试 (尝试 {attempt + 1}/{self.max_retries + 1}): {e}")
time.sleep(current_delay)
current_delay *= self.backoff_factor
else:
logger.error(f"执行失败,已达到最大重试次数: {e}")
raise last_exception