"""
霜云(Shimokumo) - 辅助函数模块
提供文本处理、时间格式化、文件操作等通用工具函数。
"""
import json
import os
import re
import uuid
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Union
def format_timestamp(
dt: Optional[datetime] = None,
fmt: str = "%Y-%m-%d %H:%M:%S",
) -> str:
"""
格式化时间戳。
Args:
dt: datetime对象,为None则使用当前时间
fmt: 格式化字符串
Returns:
格式化后的时间字符串
Example:
>>> format_timestamp()
'2026-07-06 12:00:00'
"""
if dt is None:
dt = datetime.now()
return dt.strftime(fmt)
def format_duration(seconds: float) -> str:
"""
将秒数格式化为人类可读的时长字符串。
Args:
seconds: 秒数
Returns:
格式化的时长字符串,如 "2小时30分15秒"
"""
if seconds < 0:
return "0秒"
seconds = int(seconds)
days, seconds = divmod(seconds, 86400)
hours, seconds = divmod(seconds, 3600)
minutes, seconds = divmod(seconds, 60)
parts: List[str] = []
if days > 0:
parts.append(f"{days}天")
if hours > 0:
parts.append(f"{hours}小时")
if minutes > 0:
parts.append(f"{minutes}分")
if seconds > 0 or not parts:
parts.append(f"{seconds}秒")
return "".join(parts)
def truncate_text(text: str, max_length: int = 500, suffix: str = "...") -> str:
"""
截断文本到指定长度。
Args:
text: 原始文本
max_length: 最大长度
suffix: 截断后缀
Returns:
截断后的文本
Example:
>>> truncate_text("这是一段很长的文本", max_length=5, suffix="...")
'这是一段很...'
"""
if len(text) <= max_length:
return text
return text[: max_length - len(suffix)] + suffix
def count_tokens_estimate(text: str) -> int:
"""
粗略估算文本的token数量。
对于中文,大约每个字符对应1-2个token;
对于英文,大约每4个字符对应1个token。
Args:
text: 输入文本
Returns:
估算的token数量
"""
if not text:
return 0
chinese_chars = len(re.findall(r"[\u4e00-\u9fff]", text))
other_chars = len(text) - chinese_chars
# 中文约1.5 token/字,英文约0.25 token/字符
return int(chinese_chars * 1.5 + other_chars * 0.25)
def safe_json_loads(text: str, default: Any = None) -> Any:
"""
安全地解析JSON字符串。
Args:
text: JSON字符串
default: 解析失败时返回的默认值
Returns:
解析后的Python对象,或默认值
"""
try:
return json.loads(text)
except (json.JSONDecodeError, TypeError, ValueError):
return default
def safe_json_dumps(obj: Any, indent: int = 2, ensure_ascii: bool = False) -> str:
"""
安全地将Python对象序列化为JSON字符串。
Args:
obj: 要序列化的对象
indent: 缩进空格数
ensure_ascii: 是否转义非ASCII字符
Returns:
JSON字符串
"""
try:
return json.dumps(obj, indent=indent, ensure_ascii=ensure_ascii)
except (TypeError, ValueError) as e:
return f"{{\"error\": \"JSON序列化失败: {str(e)}\"}}"
def ensure_directory(path: str) -> str:
"""
确保目录存在,不存在则创建。
Args:
path: 目录路径
Returns:
目录的绝对路径
"""
abs_path = os.path.abspath(path)
os.makedirs(abs_path, exist_ok=True)
return abs_path
def read_file_content(file_path: str, encoding: str = "utf-8") -> Optional[str]:
"""
读取文件内容。
Args:
file_path: 文件路径
encoding: 文件编码
Returns:
文件内容字符串,失败返回None
"""
try:
with open(file_path, "r", encoding=encoding) as f:
return f.read()
except (IOError, OSError) as e:
return None
def write_file_content(
file_path: str,
content: str,
encoding: str = "utf-8",
) -> bool:
"""
写入文件内容。
Args:
file_path: 文件路径
content: 要写入的内容
encoding: 文件编码
Returns:
是否写入成功
"""
try:
dir_path = os.path.dirname(file_path)
if dir_path:
os.makedirs(dir_path, exist_ok=True)
with open(file_path, "w", encoding=encoding) as f:
f.write(content)
return True
except (IOError, OSError) as e:
return False
def generate_unique_id(prefix: str = "") -> str:
"""
生成唯一ID。
Args:
prefix: ID前缀
Returns:
带前缀的唯一ID字符串
Example:
>>> generate_unique_id("chat")
'chat_3fa85f64-5717-4562-b3fc-2c963f66afa6'
"""
uid = str(uuid.uuid4())[:8]
if prefix:
return f"{prefix}_{uid}"
return uid
def clean_html_text(html: str) -> str:
"""
清理HTML标签,提取纯文本。
Args:
html: HTML字符串
Returns:
清理后的纯文本
"""
# 移除script和style标签及其内容
text = re.sub(r"", "", html, flags=re.IGNORECASE)
text = re.sub(r"", "", text, flags=re.IGNORECASE)
# 移除所有HTML标签
text = re.sub(r"<[^>]+>", "", text)
# 解码HTML实体
text = text.replace(" ", " ")
text = text.replace("<", "<")
text = text.replace(">", ">")
text = text.replace("&", "&")
text = text.replace(""", '"')
text = text.replace("'", "'")
# 合并多余空白
text = re.sub(r"\s+", " ", text).strip()
return text
def extract_urls(text: str) -> List[str]:
"""
从文本中提取URL列表。
Args:
text: 输入文本
Returns:
URL列表
"""
url_pattern = r"https?://[^\s<>\"]+|[a-zA-Z0-9-]+\.[a-zA-Z]{2,}[^\s<>\"]*"
return re.findall(url_pattern, text)
def chunk_text(
text: str,
chunk_size: int = 1000,
overlap: int = 100,
) -> List[str]:
"""
将长文本按段落分块。
Args:
text: 输入文本
chunk_size: 每块最大字符数
overlap: 块间重叠字符数
Returns:
文本块列表
"""
if len(text) <= chunk_size:
return [text]
chunks: List[str] = []
start = 0
while start < len(text):
end = start + chunk_size
# 尝试在句号、换行等处断句
if end < len(text):
for sep in ["\n", "。", ".", "!", "!", "?", "?"]:
sep_pos = text.rfind(sep, start + chunk_size // 2, end)
if sep_pos > start:
end = sep_pos + 1
break
chunks.append(text[start:end])
start = end - overlap
return chunks
def calculate_similarity(text1: str, text2: str) -> float:
"""
计算两段文本的简单相似度(基于字符重叠)。
Args:
text1: 文本1
text2: 文本2
Returns:
相似度分数 (0.0 ~ 1.0)
"""
if not text1 or not text2:
return 0.0
set1 = set(text1.lower())
set2 = set(text2.lower())
intersection = set1 & set2
union = set1 | set2
if not union:
return 1.0
return len(intersection) / len(union)
def is_media_url(url: str) -> Dict[str, Any]:
"""
检测URL是否为媒体资源URL,并识别媒体类型。
Args:
url: 要检测的URL
Returns:
包含媒体信息的字典,如 {"is_media": True, "type": "video", "format": "mp4"}
"""
media_extensions = {
"video": {
"mp4", "avi", "mkv", "mov", "wmv", "flv", "webm", "m4v", "ts", "3gp"
},
"audio": {
"mp3", "wav", "flac", "aac", "ogg", "wma", "m4a", "opus", "ape", "alac"
},
"image": {
"jpg", "jpeg", "png", "gif", "webp", "bmp", "svg", "ico", "tiff"
},
}
# 提取文件扩展名
ext_match = re.search(r"\.(\w+)(?:\?|$)", url.lower())
if not ext_match:
return {"is_media": False}
ext = ext_match.group(1)
for media_type, extensions in media_extensions.items():
if ext in extensions:
return {
"is_media": True,
"type": media_type,
"format": ext,
}
return {"is_media": False}