python Pydantic v2自定义类型验证器

5us2dqdw  于 2023-10-14  发布在  Python
关注(0)|答案(2)|浏览(290)

我正试图将我的代码更新到pydantic v2,但却找不到一个好方法来复制我在版本1中的自定义类型。我将使用自定义日期类型作为示例。最初的实现和用法看起来像这样:

from datetime import date
from pydantic import BaseModel

class CustomDate(date):
    # Override POTENTIAL_FORMATS and fill it with date format strings to match your data
    POTENTIAL_FORMATS = []
    
    @classmethod
    def __get_validators__(cls):
        yield cls.validate_date
        
    @classmethod
    def validate_date(cls, field_value, values, field, config) -> date:
        if type(field_value) is date:
            return field_value
        return to_date(field.name, field_value, cls.POTENTIAL_FORMATS, return_str=False)

class ExampleModel(BaseModel):
    class MyDate(CustomDate):
        POTENTIAL_FORMATS = ['%Y-%m-%d', '%Y/%m/%d']
    dt: MyDate

我试着按照官方文档和下面的例子here来做,它基本上是工作的,但是info参数没有我需要的字段(datafield_name)。尝试访问它们会导致AttributeError。

info.field_name
*** AttributeError: No attribute named 'field_name'

Annotated__get_pydantic_core_schema__方法都存在此问题

from datetime import date
from typing import Annotated

from pydantic import BaseModel, BeforeValidator
from pydantic_core import core_schema  

class CustomDate:
    POTENTIAL_FORMATS = []

    @classmethod
    def validate(cls, field_value, info):
        if type(field_value) is date:
            return field_value
        return to_date(info.field_name, field_value, potential_formats, return_str=False)

    @classmethod
    def __get_pydantic_core_schema__(cls, source, handler) -> core_schema.CoreSchema:
        return core_schema.general_plain_validator_function(cls.validate)

def custom_date(potential_formats):
    """
    :param potential_formats: A list of datetime format strings
    """
    def validate_date(field_value, info) -> date:
        if type(field_value) is date:
            return field_value
        return to_date(info.field_name, field_value, potential_formats, return_str=False)
    CustomDate = Annotated[date, BeforeValidator(validate_date)]
    return CustomDate

class ExampleModel(BaseModel):
    class MyDate(CustomDate):
        POTENTIAL_FORMATS = ['%Y-%m-%d', '%Y/%m/%d']
    dt: MyDate
    dt2: custom_date(['%Y-%m-%d', '%Y/%m/%d'])

如果我只是将validate_date函数作为一个常规的field_validator函数包含进来,我会得到info,其中包含了我需要的所有字段,只有在与自定义类型一起使用时,我才会看到这个问题。如何编写一个自定义类型,使其可以访问以前验证过的字段和正在验证的字段的名称?

4szc88ey

4szc88ey1#

一种更简单的方法是通过Annotated类型执行验证。但是,在某些情况下,您可能需要完全自定义的类型。

注解字段

(The简单的方法)

from datetime import datetime, date
from functools import partial
from typing import Any, List
from typing_extensions import Annotated

from pydantic import TypeAdapter
from pydantic.functional_validators import BeforeValidator

def try_parse_date(v: Any, allowed_formats: List[str]) -> Any:

    if isinstance(v, str):
        for fmt in allowed_formats:
            try:
                return datetime.strptime(v, fmt).date()
            except ValueError:
                continue
    else:
        return v

CustomDate = Annotated[
    date,
    BeforeValidator(
        partial(
            try_parse_date,
            allowed_formats=['%Y-%m-%d', '%Y/%m/%d', '%Y_%m_%d']
        )
    )
]

下面是一个测试,它确保了预期的行为:

def test_custom_type():
    values = [
        "2023-06-01", "2023/06/01", "2023_06_01",
        date(2023, 6, 1)
    ]
    expected = date(2023, 6, 1)

    ta = TypeAdapter(CustomDate)
    result = [ta.validate_python(x) for x in values]
    assert all(x==expected for x in result)

全定制类型

您遇到的问题与pydantic执行验证的顺序有关。假设date有自己的核心架构(例如:将验证时间戳或类似的转换),您将希望在核心验证之前执行验证。
相关答案(使用更简单的代码):在Pydantic v2中定义自定义类型
要解决这个问题,您需要在自定义类型中定义__get_pydantic_core_schema__。我在下面链接了模式验证,它允许将多个类型合并为一个(例如:假设你想把datetime转换成日期,你可以在链中这样做)。我还使用了general_plain_validator_function,它不需要特定的模式来操作(最普通的选择)。

Pydantic V2中创建自定义Date类型的代码:

from datetime import datetime, date
from typing import Any, List

from pydantic import BaseModel, GetCoreSchemaHandler
from pydantic_core import CoreSchema, core_schema


class CustomDate(date):
    """Custom date"""

    allowed_formats: List[str] = ['%Y-%m-%d', '%Y/%m/%d']

    @classmethod
    def try_parse_date(cls, v: Any, info: core_schema.ValidationInfo) -> Any:

        if isinstance(v, str):
            for fmt in cls.allowed_formats:
                try:
                    return datetime.strptime(v, fmt).date()
                except ValueError:
                    continue
        else:
            return v

    @classmethod
    def truncate_datetime(cls, v: Any, info: core_schema.ValidationInfo) -> Any:
        """If a datetime value is provided, truncate to a date"""
        if isinstance(v, datetime):
            return v.date()
        else:
            return v

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> CoreSchema:

        return core_schema.chain_schema(
            [
                core_schema.general_plain_validator_function(
                    function=cls.truncate_datetime,
                ),
                core_schema.general_plain_validator_function(
                    function=cls.try_parse_date,
                )

            ]
        )

您可以使用已定义的类型,也可以通过子类化该类型来自定义allow_formats

class ExampleModel(BaseModel):
    class MyDate(CustomDate):
        allowed_formats = ['%Y-%m-%d', '%Y/%m/%d', '%Y_%m_%d']

    dt: MyDate

这里有一个快速测试,显示事情正在工作:

def test_model():
    values = [
        "2023-06-01", "2023/06/01", "2023_06_01",
        date(2023, 6, 1), datetime(2023, 6, 1, 1)
    ]
    expected = date(2023, 6, 1)

    data = [ExampleModel(dt=v) for v in values]
    assert all(x.dt == expected for x in data)
l7wslrjt

l7wslrjt2#

从2.4版开始,你可以将field_name和data放在一起。点击这里查看更新的文档。
现在我的自定义数据类型的第一个版本看起来像这样:

class CustomDate:
    POTENTIAL_FORMATS = []

    @classmethod
    def validate(cls, field_value, info):
        if type(field_value) is date:
            return field_value
        return to_date(info.field_name, field_value, cls.POTENTIAL_FORMATS, return_str=False)

    @classmethod
    def __get_pydantic_core_schema__(cls, source, handler) -> core_schema.CoreSchema:
        return core_schema.with_info_before_validator_function(
            cls.validate, handler(date), field_name=handler.field_name
        )

我需要改变的只是我使用的core_schema验证器函数。第二个版本的自定义数据类型(使用Annotated的)现在可以正常工作,没有任何更改。

Pydantic 2.4之前

根据this feature request,在自定义类型验证器中访问info.datainfo.field_name目前在v2中是不可能的。
如果你只需要info.data,那么看起来你可以用core_schema.field_before_validator_function定义你的验证器(我猜所有的field_*验证器都可以工作),尽管你需要创建一个字段名:

from dataclasses import dataclass
from typing import Annotated, List, Any, Callable

from pydantic import ValidationError, BaseModel, Field, BeforeValidator, field_validator, GetCoreSchemaHandler
from pydantic_core import core_schema, CoreSchema

def fn(v: str, info: core_schema.ValidationInfo, *args, **kwargs) -> str:
    try:
        print(f'Validating {info.field_name}')
        return info.data['use_this']
    except AttributeError as err:
        return 'No data'

class AsFieldB4Method(str):
    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler, *args, **kwargs
    ) -> CoreSchema:
        return core_schema.field_before_validator_function(fn, 'not_the_real_field_name', core_schema.str_schema())

class MyModel(BaseModel):
    use_this: str
    core_schema_field_b4_method: AsFieldB4Method  # Partially works

从评论来看,听起来pydantic团队想让它与非字段验证器一起工作,并使访问info.field_name成为可能,所以希望能实现。我会更新这个答案时发生的变化,但检查该链接的情况下,我错过了它。

相关问题