在 Pydantic 中使用自定義型別

記錄一下 pydantic core 如何認識自定義型別

Pydantic 中使用非原生支援的型別會跳出錯誤訊息。舉例來說,使用時區型別的欄位會產生 PydanticSchemaGenerationError

import datetime
from pydantic import BaseModel
class User(BaseModel):
timezone: datetime.tzinfo
User(timezone=datetime.timezone.utc)
"""
PydanticSchemaGenerationError: Unable to generate pydantic-core schema for <class 'datetime.tzinfo'>...
"""

本文紀錄要怎麼使用自定義的型別。

本文撰寫於 2023 年 12 月,此時的 Pydantic 發行版本為 2.5、Pydantic Core 版本為 2.14。

InstanceOf

InstanceOf 拿出來能解決 PydanticSchemaGenerationError 的問題:

from pydantic import InstanceOf

class User(BaseModel):
    timezone: InstanceOf[datetime.tzinfo]

# ok
User(timezone=datetime.timezone.utc)

如果只需要把 Pydantic 當成進階版 dataclasses 的話那這樣就夠了。全文完,收工。

搭配 annotated validator

接下來是更複雜一點的需求:我想利用 Pydantic 來解析我從資料庫或網路拉來的資料,這些外部過來的資料大概率只會由原始型別(primitive type)組成,如 JSON 格式;而我希望能將這些資料轉成 Python 物件以便於後續操作。

那麽就可以搭配 annotated validator 來進行解析:

import zoneinfo
from typing import Annotated, Any
from pydantic import BeforeValidator
def _validate(v: Any) -> datetime.tzinfo:
if isinstance(v, str):
try:
return zoneinfo.ZoneInfo(v)
except zoneinfo.ZoneInfoNotFoundError:
raise ValueError(f"Invalid timezone: {v}")
return v
class User(BaseModel):
timezone: Annotated[InstanceOf[datetime.tzinfo], BeforeValidator(_validate)]

這裡的手段是使用了 BeforeValidator 讓其在 InstanceOf 檢查之前先把輸入的文字轉換成 ZoneInfo 物件;這樣的話以下幾個寫法就都合法了:

User(timezone=datetime.timezone.utc)
User(timezone="Asia/Taipei")
User(timezone=zoneinfo.ZoneInfo("America/New_York"))

搭配序列化

在前面我們利用了驗證器達成反序列化,那下一步就是序列化了;序列化畢竟是讓資料被儲存歸檔、或發送至其他服務前的必要過程。

這種需求可以利用序列化器如 PlainSerializer 來達成把物件轉回字串:

from pydantic import PlainSerializer
def _serialize(tz: zoneinfo.ZoneInfo) -> str:
return tz.key
class User(BaseModel):
timezone: Annotated[
InstanceOf[datetime.tzinfo],
BeforeValidator(_validate),
PlainSerializer(_serialize),
]
user = User(timezone="America/New_York")
print(user.model_dump())

有了這個序列化器,那麽這個欄位在 model_dump()model_dump_json() 的輸出就會變成字串,這樣就十分容易儲存了。

# 沒有 PlainSerializer 時
{'timezone': zoneinfo.ZoneInfo(key='America/New_York')}

# 有 PlainSerializer 時
{'timezone': 'America/New_York'}

跟驗證器類似地,PlainSerializer 的輸出會直接被使用,如果需要嵌合其他序列化函數使用,也存在 WrapSerializer 可以利用;使用 WrapSerializer 時序列化函數的函數簽章(function signature)為:

from pydantic_core.core_schema import SerializerFunctionWrapHandler
def _serialize(v: Any, serializer: SerializerFunctionWrapHandler) -> str:
"""
Parameters
----------
v : Any
要被序列化的物件;實際型別會是 `InstanceOf` 所帶的型別參數
serializer : SerializerFunctionWrapHandler
用於序列化的函數
"""

不過上面的型別註釋使用了 InstanceOf 的情境下,拿到的 serializer 會是一個 SerializationCallable(serializer=any),基本上是一個接受任何東西、輸出任何東西簡稱摸魚的序列化器;所以在上面的範例中改成 WrapSerializer 就不會感到什麼變化。

額外設定

PlainSerializerWrapSerializer 除了需要序列化函數作為參數以外,它們還接受兩個可選參數:

  • return_type - 用於提示序列化完的輸出型別,當發生型別不符時會送出警告
  • when_used - 用於指定這個序列化工具在什麼時候被採用

舉例來說,這段程式指定了僅在 json 條件下採用這個序列化函數,那麽該型別在呼叫 model_dump() 就會繼續回傳 ZoneInfo,而在 model_dump_json() 才會是以 str 回傳。

class User(BaseModel):
timezone: Annotated[
InstanceOf[datetime.tzinfo],
BeforeValidator(_validate),
PlainSerializer(_serialize, when_used="json"),
]

搭配 Pydantic core

以上的情境都是在不對目標型別做更動的情況下達成;不過如果這個型別是自定義的,那麼就可以直接實作 __get_pydantic_core_schema__ 來讓 Pydantic 認識這個型別。

實作 __get_pydantic_core_schema__

舉例來說,使用自定義的型別來支援某些東亞曆法——

import re
from typing import Self
from pydantic import GetCoreSchemaHandler
from pydantic_core import CoreSchema, core_schema
from pydantic_core.core_schema import ValidatorFunctionWrapHandler
class MinguoDate:
gregorian: datetime.date
@classmethod
def __get_pydantic_core_schema__(
cls, source_type: Any, handler: GetCoreSchemaHandler
) -> CoreSchema:
serializer = core_schema.plain_serializer_function_ser_schema(str)
inner_schema = core_schema.is_instance_schema(cls, serialization=serializer)
def _validate(v: Any):
if isinstance(v, str):
match = re.match(r"民國\s*(\d+)\s*年\s*(\d+)\s*月\s*(\d+)\s*日", v)
if match:
return cls(*map(int, match.groups()))
return v
return core_schema.no_info_before_validator_function(_validate, inner_schema)
def __init__(self, year: int, month: int, day: int):
self.gregorian = datetime.date(year + 1911, month, day)
def __repr__(self):
return f"MinguoDate({self.gregorian!r})"
def __str__(self) -> str:
return f"民國{self.gregorian.year - 1911}{self.gregorian.month}{self.gregorian.day}日"
class User(BaseModel):
birthday: MinguoDate
user = User(birthday="民國112年 12月 4日")
print(user)
"""
birthday=MinguoDate(datetime.date(2023, 12, 4))
"""

__get_pydantic_core_schema__ 是 Pydantic Core 實際上在理解各個型別時所使用的統一介面。這個介面必須為一個類別方法(class method)並回傳 core schema 以讓 Pydantic Core 得以識別與處理。

通常來說我們不需要從頭到尾去寫出整個 core schema,而是搭配 pydantic_core.core_schema 中所提供的輔助函數來建立相關的物件。以上面的程式做舉例,它包含了以下幾個設定:

補充一點,對於各種驗證器在 core schema 這層都有提供 no_info_*with_info_* 兩個版本,差別在於後者傳入驗證函數的參數會多一個 ValidationInfo 提供當下要被驗證的物件的資訊如欄位名稱及此欄位的 core schema;這個參數在要實作泛用的驗證器時可以利用,而如果僅需對應特定格式則前者就滿足需求了。

作為註釋

__get_pydantic_core_schema__ 介面不僅可以被實作在自定義的型別上、我們也可以使用它來單純提供詮釋資料(metadata)並與 Annotated 搭配使用。

例如將上面的型別改寫一下,就會變成接受民國紀元字串的原生 date 物件:

class MinguoDateFormat:
@classmethod
def __get_pydantic_core_schema__(
cls, source_type: Any, handler: GetCoreSchemaHandler
) -> CoreSchema:
inner_schema = core_schema.is_instance_schema(datetime.date)
def _validate(v: Any):
if isinstance(v, str):
match = re.match(r"民國\s*(\d+)\s*年\s*(\d+)\s*月\s*(\d+)\s*日", v)
if match:
year, month, day = match.groups()
return datetime.date(int(year) + 1911, int(month), int(day))
return v
return core_schema.no_info_before_validator_function(_validate, inner_schema)
class User(BaseModel):
birthday: Annotated[datetime.date, MinguoDateFormat]

接受作為註釋這個特性在使用於第三方函式庫時看起來似乎比較彆扭——畢竟程式碼不會像是使用 annotated validator 那麽乾淨。 但這個特性其實來自於原生的需求,Pydantic 中各種驗證器及序列化器註釋的底層實作皆是提供了這個介面來讓 Pydantic Core 得以識別與處理。

小結

  • 使用 InstanceOf 於非原生型別上
  • 可以把驗證器(validator)拿來當作反序列化的工具
  • 使用序列化函數(serializer)將物件序列化以用於儲存