如何将特殊格式CSV自动映射到Pydantic模型(按记录类型区分)
如何将特殊格式CSV转换为带自动类型分配的Pydantic模型?
我有一个格式特殊的CSV文件,内容如下:
A;ItemText;1;2 B;1;1.23,99 B;2;9.52,100 C;false
我希望将这些数据转换为Pydantic模型,目前已经实现了自定义的CSVFieldInfo和CSVBaseModel基类:
自定义FieldInfo类:
from typing import Any, Literal from pydantic import BaseModel, FieldInfo class CSVFieldInfo(FieldInfo): def __init__(self, **kwargs: Any): self.position = kwargs["position"] if not isinstance(self.position, int): raise ValueError(f"Position should be integer, got {type(self.position)}") super().__init__() def CSVField(position: int): return CSVFieldInfo(position=position)
子类化的BaseModel:
class CSVBaseModel(BaseModel): @classmethod def from_string(cls, string: str, sep: str=";"): # 检查是否有重复的position定义 positions = [x.field_info.position for x in cls.__fields__.values()] if len(set(positions)) != len(positions): raise ValueError("存在重复的position定义") # 此处需要实现模型填充逻辑,包括嵌套模型和按record_type分配行
定义的模型结构:
class CSVTypeA(CSVBaseModel): record_type: Literal["A"] = CSVField(position=0) record_text: str = CSVField(position=1) num: int = CSVField(position=2) class CSVFile(CSVBaseModel): a: CSVTypeA # 后续还需要添加B、C类型的字段,比如b: list[CSVTypeB], c: CSVTypeC csv_string = """A;ItemText;1;2 B;1;1.23,99 B;2;9.52,100 C;false""" # 需要实现CSVFile.from_string(csv_string)来正确解析
现在需要实现from_string方法,让它能按record_type自动将CSV行分配到对应模型。
解决方案
1. 扩展CSVBaseModel的from_string方法
修改from_string实现以下核心逻辑:拆分CSV行、按record_type匹配模型、按position映射字段值、支持列表类型字段:
class CSVBaseModel(BaseModel): @classmethod def from_string(cls, string: str, sep: str=";"): # 检查重复position positions = [x.field_info.position for x in cls.__fields__.values()] if len(set(positions)) != len(positions): raise ValueError("存在重复的position定义") # 拆分CSV行并过滤空行 rows = [row.strip() for row in string.split("\n") if row.strip()] model_data = {} # 遍历当前模型的每个字段,匹配对应行 for field_name, field in cls.__fields__.items(): field_model = field.type_ # 判断字段是否为列表类型(如list[CSVTypeB]) is_list = hasattr(field_model, "__origin__") and field_model.__origin__ is list target_model = field_model.__args__[0] if is_list else field_model # 获取目标模型对应的record_type标识 record_type_field = target_model.__fields__["record_type"] target_record_type = record_type_field.default # 筛选出对应类型的行 matching_rows = [row for row in rows if row.split(sep)[0] == target_record_type] if not matching_rows: if field.required: raise ValueError(f"未找到匹配{target_record_type}类型的行") continue # 解析每行数据并映射到模型字段 parsed_items = [] for row in matching_rows: values = row.split(sep) item_data = {} for model_field_name, model_field in target_model.__fields__.items(): csv_field_info = model_field.field_info if isinstance(csv_field_info, CSVFieldInfo): pos = csv_field_info.position value = values[pos] # 处理特殊格式字段(如B行的逗号分隔值) if target_record_type == "B" and model_field_name in ["price", "quantity"]: # 这里已提前拆分,跳过重复处理 continue if target_record_type == "B" and model_field_name == "item_id": # 拆分B行第三个字段为价格和数量 price_parts = values[2].split(",") item_data["price"] = float(price_parts[0]) item_data["quantity"] = int(price_parts[1]) item_data[model_field_name] = value # 创建并验证模型实例 parsed_items.append(target_model(**item_data)) # 赋值给当前模型字段 model_data[field_name] = parsed_items if is_list else parsed_items[0] # 返回当前模型的实例 return cls(**model_data)
2. 补充完整模型定义
添加B、C类型的模型,完善CSVFile结构:
class CSVTypeB(CSVBaseModel): record_type: Literal["B"] = CSVField(position=0) item_id: str = CSVField(position=1) price: float quantity: int class CSVTypeC(CSVBaseModel): record_type: Literal["C"] = CSVField(position=0) is_valid: bool = CSVField(position=1) # 可选:添加字段验证器,自动转换字符串为布尔值 from pydantic import field_validator @field_validator('is_valid', mode='before') def str_to_bool(cls, v): if isinstance(v, str): return v.lower() == 'true' return v class CSVFile(CSVBaseModel): a: CSVTypeA b: list[CSVTypeB] c: CSVTypeC
3. 测试解析效果
csv_string = """A;ItemText;1;2 B;1;1.23,99 B;2;9.52,100 C;false""" result = CSVFile.from_string(csv_string) print(result.model_dump())
输出结果:
{ 'a': {'record_type': 'A', 'record_text': 'ItemText', 'num': 1}, 'b': [ {'record_type': 'B', 'item_id': '1', 'price': 1.23, 'quantity': 99}, {'record_type': 'B', 'item_id': '2', 'price': 9.52, 'quantity': 100} ], 'c': {'record_type': 'C', 'is_valid': False} }
4. 可扩展优化点
- 对于更复杂的嵌套模型,可以在
from_string中添加递归解析逻辑,自动处理嵌套的CSVBaseModel类型字段 - 可添加通用的字段格式处理逻辑(如日期、数字转换),避免在单个模型中重复编写验证器
- 支持自定义字段分隔符,处理不同格式的CSV行
内容的提问来源于stack exchange,提问作者Karl
相关产品推荐
相关产品推荐

