Pydantic 2.0最新版 - 【自定义数据类型详解，史诗级详细解释，深入学习】

#!/usr/bin/python
# -*- coding: utf-8 -*-
# @version        : 1.0
# @Create Time    : 2023/7/15 21:45
# @File           : pydantic.py
# @IDE            : PyCharm
# @desc           : pydantic 2.0 自定义数据类型方法

print("请大家支持一下我的开源项目：https://gitee.com/ktianc/kinit")
# 官方文档：https://docs.pydantic.dev/dev-v2/usage/types/custom/#adding-validation-and-serialization

import datetime
from typing import Annotated
from pydantic import BaseModel, AfterValidator, PlainSerializer, WithJsonSchema


def after(value: str | datetime.datetime | int | float):
    """
    执行逻辑：
    当值传入后立即执行
    所以在使用 .属性 或者 序列化方法时，都是获取到的处理后的数据
    """
    print("after 方法执行了", value)
    print(type(value))
    if isinstance(value, str):
        pattern = "%Y-%m-%d %H:%M:%S"
        try:
            datetime.datetime.strptime(value, pattern)
            return value
        except ValueError:
            pass
    elif isinstance(value, datetime.datetime):
        return value.strftime("%Y-%m-%d %H:%M:%S")
    raise ValueError("无效的日期时间或字符串数据")


def plain(value: str):
    """
    执行逻辑：
    只有在使用 model_dump 等序列化方式时，才会执行，使用 .属性 时并不会执行
    所以如果有数据需要在使用 .属性 或者 序列化时 处理的结果一样，那么请在 after 方法中处理
    如果 after 方法与 plain 方法处理的是一样的，那么可以在这里直接返回值，或者使用 lambda x: x 匿名函数直接返回
    """
    print("plain 方法执行了", value)
    return value


# 自定义数据类型
# 实现自定义一个日期时间字符串的数据类型
# 如果我传入的是字符串，那么直接返回，如果我传入的是一个日期类型，那么会转为字符串格式后返回
# 因为在 pydantic 2.0 中是支持 int 或 float 自动转换类型的，所以我这里添加进去，但是在处理时会使这两种类型报错
# 官方文档：https://docs.pydantic.dev/dev-v2/usage/types/datetime/
TruncatedDatetimeStr = Annotated[
    str | datetime.datetime | int | float,
    AfterValidator(after),
    PlainSerializer(plain, return_type=str),
    WithJsonSchema({'type': 'string'}, mode='serialization')
]


class Model(BaseModel):
    """
    使用DateTime自定义数据类型的Pydantic模型示例。
    """

    date: TruncatedDatetimeStr


if __name__ == '__main__':
    # 测试使用字符串输入
    m1 = Model(date="2023-07-15 21:34:56")
    print(m1.date)  # 输出: 2023-07-15 21:34:56
    print(type(m1.date))  # 输出: <class 'str'>
    print(m1.model_dump())  # 输出：{'date': '2023-07-15 21:34:56'}
    print(type(m1.model_dump()['date']))  # 输出：<class 'str'>

    print("\n===================================================\n")

    # 测试使用datetime输入
    m2 = Model(date=datetime.datetime.now())
    print(m2.date)  # 输出: 当前日期时间字符串
    print(type(m2.date))  # 输出: <class 'str'>
    print(m2.model_dump())  # 输出：{'date': 当前日期时间字符串}
    print(type(m2.model_dump()['date']))  # 输出：<class 'str'>

    print("\n===================================================\n")

    # 测试无效的日期时间字符串
    try:
        m3 = Model(date="2023-07-15 12:34:61")  # 这里的秒，填的是错误的
    except ValueError as exc:
        print(exc)  # 输出: 无效的日期时间字符串

    print("\n===================================================\n")

    # 测试无效的输入类型
    try:
        m4 = Model(date=123)
        print(m4.date)
    except ValueError as exc:
        print(exc)  # 输出: 无效的日期时间字符串

    print("\n===================================================\n")