136 lines
4.9 KiB
Python
136 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
文件输入输出工具模块
|
||
"""
|
||
|
||
import os
|
||
import random
|
||
import json
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Optional, List, Dict, Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ResourceLoader:
|
||
"""资源加载器，用于加载文件内容"""
|
||
|
||
@staticmethod
|
||
def load_text_file(file_path: str) -> Optional[str]:
|
||
"""加载文本文件内容"""
|
||
if not os.path.exists(file_path):
|
||
logger.warning(f"文件不存在: {file_path}")
|
||
return None
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
return f.read()
|
||
except Exception as e:
|
||
logger.error(f"加载文件 '{file_path}' 失败: {e}")
|
||
return None
|
||
|
||
@staticmethod
|
||
def load_json_file(file_path: str) -> Optional[Dict[str, Any]]:
|
||
"""加载并解析JSON文件"""
|
||
content = ResourceLoader.load_text_file(file_path)
|
||
if content is None:
|
||
return None
|
||
try:
|
||
return json.loads(content)
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"解析JSON文件 '{file_path}' 失败: {e}")
|
||
return None
|
||
|
||
@staticmethod
|
||
def find_file(directory: str, file_name: str, exact_match: bool = True) -> Optional[str]:
|
||
"""在目录中查找文件，支持精确和模糊匹配"""
|
||
if not os.path.isdir(directory):
|
||
logger.warning(f"目录不存在: {directory}")
|
||
return None
|
||
|
||
# 确保文件名有.txt后缀，如果需要
|
||
base, ext = os.path.splitext(file_name)
|
||
if not ext:
|
||
file_name_with_ext = f"{base}.txt"
|
||
else:
|
||
file_name_with_ext = file_name
|
||
|
||
# 精确匹配
|
||
exact_path = os.path.join(directory, file_name_with_ext)
|
||
if os.path.exists(exact_path):
|
||
return exact_path
|
||
|
||
# 模糊匹配
|
||
if not exact_match:
|
||
for f in os.listdir(directory):
|
||
if base in f:
|
||
return os.path.join(directory, f)
|
||
|
||
logger.warning(f"在 '{directory}' 中找不到文件 '{file_name}'")
|
||
return None
|
||
|
||
|
||
class OutputManager:
|
||
"""负责处理输出文件，如保存文章和生成汇总报告"""
|
||
|
||
def __init__(self, output_dir: str, run_id: str):
|
||
self.base_output_dir = Path(output_dir)
|
||
self.run_id = run_id
|
||
self.run_dir = self.base_output_dir / self.run_id
|
||
self.run_dir.mkdir(parents=True, exist_ok=True)
|
||
logger.info(f"OutputManager initialized for run '{run_id}' in '{self.run_dir}'")
|
||
|
||
def get_topic_dir(self, topic_index: Any) -> Path:
|
||
"""为给定主题索引创建并返回一个唯一的目录"""
|
||
topic_dir = self.run_dir / f"topic_{topic_index}"
|
||
topic_dir.mkdir(parents=True, exist_ok=True)
|
||
return topic_dir
|
||
|
||
def get_variant_dir(self, topic_index: int, variant_index: int) -> Path:
|
||
"""获取并创建特定变体的目录"""
|
||
variant_dir = self.run_dir / f"{topic_index}_{variant_index}"
|
||
variant_dir.mkdir(exist_ok=True)
|
||
return variant_dir
|
||
|
||
def save_json(self, data: Any, filename: str, subdir: Optional[str] = None):
|
||
"""将数据保存为JSON文件"""
|
||
target_dir = self.run_dir / subdir if subdir else self.run_dir
|
||
target_dir.mkdir(exist_ok=True)
|
||
file_path = target_dir / filename
|
||
try:
|
||
with open(file_path, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=4)
|
||
logger.info(f"JSON data saved to: {file_path}")
|
||
except Exception as e:
|
||
logger.error(f"Failed to save JSON to {file_path}: {e}")
|
||
|
||
def save_text(self, content: str, filename: str, subdir: Optional[str] = None):
|
||
"""将文本内容保存为文件"""
|
||
target_dir = self.run_dir / subdir if subdir else self.run_dir
|
||
target_dir.mkdir(exist_ok=True)
|
||
file_path = target_dir / filename
|
||
try:
|
||
with open(file_path, 'w', encoding='utf-8') as f:
|
||
f.write(content)
|
||
logger.info(f"Text data saved to: {file_path}")
|
||
except Exception as e:
|
||
logger.error(f"Failed to save text to {file_path}: {e}")
|
||
|
||
def save_image(self, image_data, filename: str, subdir: Optional[str] = None):
|
||
"""保存图像文件 (需要Pillow库)"""
|
||
target_dir = self.run_dir / subdir if subdir else self.run_dir
|
||
target_dir.mkdir(exist_ok=True)
|
||
file_path = target_dir / filename
|
||
try:
|
||
image_data.save(file_path)
|
||
logger.info(f"Image saved to: {file_path}")
|
||
except Exception as e:
|
||
logger.error(f"Failed to save image to {file_path}: {e}")
|
||
|
||
def finalize(self):
|
||
"""完成运行的最终操作"""
|
||
logger.info(f"Finalizing run: {self.run_id}")
|
||
# 目前没有特殊操作，但可以用于未来的扩展，如创建清单文件
|
||
pass |