在 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
就不會感到什麼變化。
額外設定
PlainSerializer
及 WrapSerializer
除了需要序列化函數作為參數以外,它們還接受兩個可選參數:
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
中所提供的輔助函數來建立相關的物件。以上面的程式做舉例,它包含了以下幾個設定:
- 最中心的
is_instance_schema
,提供等價於InstanceOf
的 core schema, 用於確認物件需要是一個MinguoDate
型別 - 在
is_instance_schema
內帶了一個plain_serializer_function_ser_schema
,對應於PlainSerializer
,用來聲明這個物件在做序列化時可以直接使用__str__
- 外掛的
no_info_before_validator_function
,對應於BeforeValidator
,提供額外的反序列化方法
補充一點,對於各種驗證器在 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)將物件序列化以用於儲存