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