TravelContentCreator/TWEET/utils/file_io.py

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