InterfaceAutoTest/commons/file_processors/yaml_processor.py

#!/usr/bin/env python
# coding=utf-8

"""
@author: CNWei
@Software: PyCharm
@contact: t6i888@163.com
@file: yaml_processor
@date: 2025/3/4 17:28
@desc: 
"""
import logging
from typing import Union, Any
from pathlib import Path
import yaml
from commons.file_processors.base_processor import BaseFileProcessor



logger = logging.getLogger(__name__)


class YamlLoadError(Exception):
    """自定义 YAML 加载异常：当 YAML 语法错误或不符合业务结构时抛出"""
    pass


class YamlProcessor(BaseFileProcessor):
    """
    用于处理 YAML 文件的类，继承自 dict。
    提供了从文件加载、保存到文件、转换为字符串和从字符串转换的功能,
    并可以直接像字典一样访问 YAML 数据。
    """

    def __init__(self, filepath: Union[str, Path], data: Union[dict, None] = None):
        """
        初始化 YamlFile 对象。

        Args:
            filepath: YAML 文件的路径 (可以是字符串或 pathlib.Path 对象).
            data: 可选的初始数据字典。如果提供，则用该字典初始化 YamlFile。
                  如果不提供，则尝试从 filepath 加载数据。
        """

        super().__init__(filepath=filepath)

        self.filepath: Path = Path(filepath)  # 确保 filepath 是 Path 对象

    def load(self) -> dict[str, Any]:
        """
        加载 YAML 文件并返回字典。

        Returns:
            Dict: 加载后的数据字典。

        Raises:
            YamlLoadError: 文件读取或解析过程中出现异常。
        """
        if not self.filepath.exists():
            logger.error(f"❌ 文件未找到: {self.filepath}")
            return {}
        try:
            with open(self.filepath, "r", encoding="utf-8") as f:
                content = yaml.safe_load(f)
                # 情况1：文件内容为空
                if content is None:
                    return {}
                    # 情况2：YAML 语法正确但不是字典（如单纯的字符串或列表）
                if not isinstance(content, dict):
                    raise YamlLoadError(f"YAML 顶层格式错误：期望 dict，实际为 {type(content).__name__}")

                return content
        except yaml.YAMLError as e:
            msg = f"❌ YAML 语法错误 [{self.filepath.name}]: {e}"
            logger.error(msg)
            raise YamlLoadError(msg) from e
        except Exception as e:
            logger.error(f"📂 读取文件系统异常: {e}")
            raise

    @staticmethod
    def to_string(data: dict[str, Any]) -> str:
        """
        将字典 (自身) 转换为 YAML 格式的字符串。

        Returns:
            YAML 格式的字符串。
        """
        try:
            return yaml.safe_dump(
                data,
                allow_unicode=True,
                sort_keys=False,
                default_flow_style=False
            )
        except TypeError as e:
            logger.error(f"将数据转换为 YAML 字符串时出错: {e}")
            return ""
        except Exception as e:
            logger.error(f"序列化 YAML 失败: {e}")
            return ""

    @staticmethod
    def from_string(yaml_str: str) -> Union[None, dict]:
        """
        将 YAML 格式的字符串转换为字典，并更新当前字典的内容.

        Args:
            yaml_str: YAML 格式的字符串。
        """
        try:
            data = yaml.safe_load(yaml_str)
            return data if isinstance(data, dict) else {}
        except yaml.YAMLError as e:
            logger.error(f"YAML 字符串解析失败: {e}")
            return {}

    def save(self, data: dict[str, Any], new_filepath: Union[str, Path, None] = None):
        """
        将字典数据保存为 YAML 文件。

        Args:
            data: 要保存的字典数据。
            new_filepath: 可选，保存到新路径。
        """
        target_path = Path(new_filepath) if new_filepath else self.filepath
        try:
            target_path.parent.mkdir(parents=True, exist_ok=True)
            with open(target_path, "w", encoding="utf-8") as f:
                yaml.safe_dump(
                    data,
                    stream=f,
                    allow_unicode=True,
                    sort_keys=False,
                    default_flow_style=False
                )
            logger.debug(f"💾 数据已成功保存至: {target_path}")
        except Exception as e:
            logger.error(f"🚫 保存 YAML 失败: {e}")
            raise
        except (TypeError, OSError) as e:
            logger.error(f"保存 YAML 文件 {self.filepath} 时出错: {e}")

# todo 需要将异常的情况返回给上层而不是默认处理为{}

if __name__ == '__main__':
    from core.settings import TEST_CASE_DIR
    # 示例用法
    yaml_path = TEST_CASE_DIR / r'answer/test_1_status.yaml'  # 你的 YAML 文件路径
    yaml_file = YamlProcessor(yaml_path)
    print(yaml_file.load())
    print(yaml_file.to_string(yaml_file.load()))
    print(type(yaml_file))

    # # 直接像字典一样访问数据
    # print("加载的数据:", yaml_file)  # 直接打印对象，就是打印字典内容
    # print("title:", yaml_file.get("title"))  # 使用 get 方法
    # if "title" in yaml_file:  # 使用 in 检查键
    #     print("原始title:", yaml_file["title"])  # 使用方括号访问
    #     yaml_file["title"] = "新的标题"  # 使用方括号修改
    #     print("修改后的title:", yaml_file["title"])
    # #
    # yaml_file["new_key"] = "new_value"  # 添加新的键值对
    #
    # # 将字典转换为 YAML 字符串
    # yaml_string = yaml_file.to_string()
    # print("\nYAML 字符串:", yaml_string)
    # #
    # # 将 YAML 字符串转换回字典 (并更新 yaml_file)
    # yaml_file.to_dict(yaml_string)
    # print("\n从字符串加载的数据:", yaml_file)
    #
    # # 保存修改后的数据 (覆盖原文件)
    # yaml_file.save()
    #
    # # 保存到新文件
    # new_yaml_path = r'D:\CNWei\CNW\InterfaceAutoTest\TestCases\test_1_user_new.yaml'
    # yaml_file.save(new_filepath=new_yaml_path)

    # 测试从字符串初始化
    # yaml_string2 = """
    # name: Test User
    # age: 30
    # """

    # yaml_file2 = YamlFile("test2.yaml", data=yaml.safe_load(yaml_string2))  # 从字符串初始化
    # print("\n从字符串初始化的 YamlFile:", yaml_file2)
    # yaml_file2.save()  # 保存到 test2.yaml
    #
    # 测试文件不存在的情形
    # non_existent_file = YamlFile("non_existent_file.yaml")
    # print("\n加载不存在的文件:", non_existent_file)  # 应该打印空字典 {}
    # non_existent_file['a'] = 1  # 可以直接添加
    # print("\n加载不存在的文件:", non_existent_file)