如何用Pydantic/SQLModel解析CSV导出的多对多关系字符串

问题描述 #

从FastAPI/SQLModel/SQLAdmin后端导出CSV时，多对多关系字段assessed_in被导出为Period模型实例的字符串形式，示例如下：

"[Period(id=1, name='1. quarter 2022', period_id='q1_2022', status=<PeriodStatusEnum.in_assessment: 'in_assessment'>), Period(id=2, name='2. Quarter 2022', period_id='q2_2022', status=<PeriodStatusEnum.in_assessment: 'in_assessment'>)]"

读取CSV后，其他字段可正常解析，但assessed_in因格式问题无法通过Dimension.model_validate完成解析。数据库模型定义：

# Database Models

class PeriodStatusEnum(str, enum.Enum):
    in_assessment = "in_assessment"
    readonly = "readonly"
    hidden = "hidden"

class Period(BaseClass, table=True):
    period_id: str = Field(index=True, unique=True)
    name: str
    status: PeriodStatusEnum = Field(sa_column=Column(Enum(PeriodStatusEnum)))

class DimensionPeriodLink(SQLModel, table=True):
    dimension_id: int | None = Field(default=None, foreign_key="dimension.id", primary_key=True)
    period_id: int | None = Field(default=None, foreign_key="period.id", primary_key=True)
    
class Dimension(BaseClass, table=True):
    dimension_id: str = Field(index=True, unique=True)
    name_EN: str
    name_DE: str
    order_nr: int
    assessed_in: list[Period] = Relationship(link_model=DimensionPeriodLink)

读取CSV得到的字典示例：

dimension_dict = {'assessed_in': "[Period(status=<PeriodStatusEnum.in_assessment: 'in_assessment'>, id=1, name='1. Quartal 2022', period_id='q1_2022'), Period(status=<PeriodStatusEnum.in_assessment: 'in_assessment'>, id=2, name='2. Quartal 2022', period_id='q2_2022')]", 'level_definitions': '[]', 'questions': '[]', 'id': '1', 'dimension_id': 'DIM10', 'name_EN': 'Organizational structure and roles', 'name_DE': 'Organisationsstruktur und Rollen', 'order_nr': '1'}

解决方案 #

方法1：自定义字符串解析函数，转换为可验证的Period列表 #

针对导出的特殊字符串格式，编写解析函数提取Period属性，转成字典后创建实例：

import re
from enum import Enum
from sqlmodel import SQLModel, Field, Column, Enum as SQLModelEnum

# 确保PeriodStatusEnum和Period模型已定义

def parse_assessed_in_str(assessed_in_str: str) -> list[Period]:
    cleaned_str = assessed_in_str.strip('"[]')
    if not cleaned_str:
        return []
    
    # 分割多个Period实例字符串
    period_strings = re.split(r'),\s*Period\(', cleaned_str)
    period_strings[0] = period_strings[0].replace('Period(', '')
    period_strings[-1] = period_strings[-1].rstrip(')')
    
    periods = []
    for p_str in period_strings:
        attrs = {}
        # 匹配key=value格式，处理字符串和枚举值
        matches = re.findall(r'(\w+)=([^\s,]+|\'.*?\'|<.*?>)', p_str)
        for key, value in matches:
            if key == 'status':
                # 提取枚举的实际值
                enum_value = re.search(r"'(\w+)'", value).group(1)
                attrs[key] = PeriodStatusEnum(enum_value)
            elif value.startswith("'") and value.endswith("'"):
                attrs[key] = value.strip("'")
            elif value.isdigit():
                attrs[key] = int(value)
            else:
                attrs[key] = value
        periods.append(Period(**attrs))
    return periods

# 使用示例
dimension_dict['assessed_in'] = parse_assessed_in_str(dimension_dict['assessed_in'])
new_record = Dimension.model_validate(dimension_dict)

方法2：修改SQLAdmin导出逻辑，输出结构化JSON #

如果能修改SQLAdmin配置，自定义assessed_in字段的导出格式，直接输出JSON数组：

from sqladmin import ModelView
from sqlmodel import Session
import json

class DimensionAdmin(ModelView, model=Dimension):
    def format_assessed_in(self, obj: Dimension, session: Session) -> str:
        # 将Period列表转为JSON字符串
        periods = [{"id": p.id, "period_id": p.period_id, "name": p.name, "status": p.status.value} for p in obj.assessed_in]
        return json.dumps(periods)
    
    # 覆盖导出字段的格式化规则
    column_formatters_export = {
        "assessed_in": format_assessed_in,
    }

导出后读取CSV时，直接解析JSON即可：

import json

dimension_dict['assessed_in'] = [Period(**p) for p in json.loads(dimension_dict['assessed_in'])]
new_record = Dimension.model_validate(dimension_dict)

方法3：自定义Pydantic字段类型自动解析 #

创建自定义字段类型，让Pydantic自动处理特殊字符串格式：

from pydantic import GetJsonSchemaHandler
from pydantic_core import core_schema
from sqlmodel import SQLModel

# 复用方法1的解析函数
def parse_assessed_in_str(assessed_in_str: str) -> list[Period]:
    # 实现同方法1
    pass

class PeriodListField(list[Period]):
    @classmethod
    def __get_pydantic_core_schema__(cls, source_type, handler: GetJsonSchemaHandler) -> core_schema.CoreSchema:
        def parse_str(value: str) -> list[Period]:
            return parse_assessed_in_str(value)
        
        return core_schema.no_info_wrap_validator_function(
            parse_str,
            core_schema.list_schema(handler(Period)),
            serialization=core_schema.plain_serializer_function_ser(lambda x: str(x)),
        )

# 创建专门用于CSV解析的验证模型
class DimensionCSVCreate(SQLModel):
    dimension_id: str
    name_EN: str
    name_DE: str
    order_nr: int
    assessed_in: PeriodListField
    # 其他字段...

# 使用示例
new_record = DimensionCSVCreate.model_validate(dimension_dict)
db_record = Dimension.from_orm(new_record)

内容的提问来源于stack exchange，提问作者Simon