You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

如何将特殊格式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

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.07.27 07:34:56