Nov 26, 2023

在 Pydantic 使用 annotated validators

Annotated validator 算是 Pydantic 魔法的延伸，它讓我們用註釋去完成資料驗證

本文撰寫於 2023 年 11 月，此時的 Pydantic 發行版本為 2.5。

Annotated validator 算是 Pydantic 魔法的延伸，過去我們利用型別提示（type hint; PEP-484）設定欄位型別檢查，而 annotated validator 讓我們用註釋（annotation; PEP-526)）設定驗證器。

驗證器顧名思義是使用於驗證欄位中的值是否符合要求，不過在接下來的段落中我們會發現可以利用它來修改輸入，進而拿來當作反序列化的手段。

標準的 annotated validator 有 AfterValidator, BeforeValidator, PlainValidator 及 WrapValidator 四種。分別對應的場合為：

AfterValidator：對 Pydantic 已經做完正規化的值做操作
BeforeValidator：直接對原始輸入值做操作、再交給 Pydnatic 去確保型別
PlainValidator：完全由這個驗證器說的算
WrapValidator：跟其他的驗證器協作

以下逐步說明使用情境。

使用後驗證器（`AfterValidator`）

這裡設計一個場景方便後續說明——以 Shop 物件去紀錄店舖的參數：

from pydantic import BaseModel

class Shop(BaseModel):
    name: str

然後我想加上一些簡單粗暴的手段擋住一些敏感詞，那麼 AfterValidator 就會是我的好夥伴：

import re
from typing import Annotated
from pydantic import AfterValidator

def _validator(v: str):
    if re.search(r"<\s*script", v):
        raise AssertionError("Please do not XSS me")
    return v

class Shop(BaseModel):
    name: Annotated[str, AfterValidator(_validator)]

# ok
Shop(name='Demo')

# error
Shop(name='<script>alert();</script>')
"""
ValidationError: 1 validation error for Shop
name
  Assertion failed, Please do not XSS me [type=assertion_error, input_value='<script>alert();</script>', input_type=str]
    For further information visit https://errors.pydantic.dev/2.5/v/assertion_error
"""

補充

Annotated validator 應用上皆需要使用 Annotated，這個型別在 Python 3.9 才加入到 typing 函式庫中，較早期的 Python 版本中則需要從 typing-extensions 中取得。

在遇到不合格的資料時我們有幾個選項：

I. ValueError, AssertionError 及 PydanticCustomError

驗證函數送出的 ValueError, AssertionError 或 PydanticCustomError 會被 Pydantic 攔截並成為 ValidationError 中錯誤訊息的一部分，如上方範例所陳列。

由於 assert 斷言檢查失敗會送出 AssertionError，所以上面的範例等價於這個寫法：

def _validator(v: str):
    assert not re.search(r"<\s*script", v), "Please do not XSS me"
    return v

但需要注意在生產環境使用 assert 是有可能被繞過的，因此個人不推薦直接使用斷言語法。

II. 其他例外型別

基本上就是直接炸，沒懸念的那種：

def _validator(v: str):
    if re.search(r"<\s*script", v):
        raise RuntimeError("Please do not XSS me")
    return v

"""
RuntimeError: Please do not XSS me
"""

III. 河蟹

我們可以在函數內修改資料，所以對部分可以忍受的錯誤可以採取河蟹手段。例如將那些 HTML 直接跳脫掉：

import html

def _validator(v: str):
    return html.escape(v)

print(Shop(name='<script>alert();</script>'))
"""
name='&lt;script&gt;alert();&lt;/script&gt;'
"""

小結，AfterValidator 是一個十分接近資料驗證這個需求的工具，並且由於在資料輸入函數時型別已經確認，我們可以更專注在「驗證」這項工作上。

使用前驗證器（`BeforeValidator`）

接下來幫 Shop 加上新的欄位：

import datetime

class Shop(BaseModel):
    created_at: datetime.datetime

Pydantic 對於 datetime.datetime 物件支援使用：

ISO 8601 格式字串

YYYY-MM-DD[T]HH:MM[:SS[.ffffff]][Z or [±]HH[:]MM]

Unix 時間

即下列兩種寫法建立出來的物件是完全相同的：

Shop(created_at=1700841600)
Shop(created_at='2023-11-25T00:00+08:00')

但我們也許需要台灣最常見的 YYYY/MM/DD，或是對岸那混亂邪惡 MM/DD/YYYY，這時候就需要前驗證器 BeforeValidator 出馬。

前驗證器會接收到尚未處理的原始資料，並在其之後交還給 Pydantic 內建的正規化函數接手。我們可以在這個階段直接對原始資料做型別轉換：

from typing import Any
from pydantic import BeforeValidator

def _validator(v: Any) -> datetime.datetime:
    if re.fullmatch(r"\d{4}/\d{2}/\d{2}", v):
        return datetime.datetime.strptime(v, "%Y/%m/%d")
    return v

class Shop(BaseModel):
    created_at: Annotated[datetime.datetime, BeforeValidator(_validator)]

代入這個自定義的驗證器之後 YYYY/MM/DD 被接受了：

# ok
Shop(created_at='2023/11/25')

# error!?
Shop(created_at=1700841600)

但少了 Pydantic 正規化的保護，本來能過的輸入卻發生了錯誤：int 輸入被餵到了 re.fullmatch 該放字串的位置。小修改一下驗證函數就能解決：

def _validator(v: Any) -> datetime.datetime:
    if isinstance(v, str):
        if re.fullmatch(r"\d{4}/\d{2}/\d{2}", v):
            return datetime.datetime.strptime(v, "%Y/%m/%d")
    return v

因此，使用 BeforeValidator 時使用者必須小心處理任何輸入的可能性。

使用 `PlainValidator`

PlainValidator 跟 BeforeValidator 很像，但 Pydantic 不會對其輸出做其他處理，故在使用它之後有可能會發生欄位實際的資料型別跟聲明不符、並且我們也無法利用那些很方便正規化函數：

from pydantic import PlainValidator

class Shop(BaseModel):
    created_at: Annotated[datetime.datetime, PlainValidator(_validator)]

print(type(shop.created_at))
"""
<class 'int'>
"""

與 BeforeValidator 相同的是，使用者依然必須小心處理任何輸入的可能性。不同的是，使用者還得自己搞輸出的處理。

使用 `WrapValidator`

WrapValidator 顧名思義用於將其他驗證器包覆，像洋蔥一樣一層層的去建構驗證的機制；在所套用的驗證函數中，使用者可以任的選擇驗證函數需要被使用的時機；基本上就是我全都要的概念。

比方說，只接受以 Unix 時間來表示的瞬時時間：

from pydantic import WrapValidator
from pydantic_core.core_schema import ValidatorFunctionWrapHandler

def _validator(v: Any, validator: ValidatorFunctionWrapHandler) -> Any:
    if isinstance(v, int):
        return validator(v)
    raise ValueError(f"Invalid input {type(v)}")

class Shop(BaseModel):
    created_at: Annotated[datetime.datetime, WrapValidator(_validator)]

在上面的範例中，驗證函數 _validator 限制了僅在輸入值為整數時才會將輸入送給內層驗證器做處理；而範例中的「內層」會是 Pydantic 為 datetime 提供的正規化函數。

另外由於驗證的時間點已經可以控制了，如果要貪心連資料格式驗證都整進來也是可以的。

由於 WrapValidator 會將內層驗證器直接交給使用者去決定操作的時機，在不當應用時它與 PlainValidator 一樣可以導致欄位的型別跟宣告的不同。

結語

標準的資料驗證場景使用 AfterValidator 會比較省力。
需要額外支援特定格式時可以利用 BeforeValidator、但要注意輸入的資料型別會是 Any。
PlainValidator 想不太到優勢
WrapValidator 可以彈性控制內層驗證器被使用的時間點——同時可以做到 BeforeValidator 跟 AfterValidator 能做到的事情。

在 Pydantic 使用 annotated validators

Annotated validator 算是 Pydantic 魔法的延伸，它讓我們用註釋去完成資料驗證

使用後驗證器（AfterValidator）