🐍

PydanticでEmailStrを拡張する

2024/10/17に公開

始めに

pydanticにはEmailStrというemailを検証するための拡張クラスがあります。しかし、Emailの仕様としてはUTF-8を許容しているものの、システム的にはASCIIしか許容したくないことがあります。その場合に向けて、EmailStrを継承してASCIIのみ許容する拡張クラスを作ります。

環境

  • Python
    • 3.12.7
  • Pydantic
    • 2.9.2

実装

まずは、ASCII文字だけチェックできるように正規表現を用意します。

^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$

次のコードでEmailStrを継承しつつ、自分が追加したvalidationを実行させます。

from typing import Type, Any
import re

from pydantic import EmailStr, validator, GetCoreSchemaHandler
from pydantic_core import core_schema


class CustomEmailStr(EmailStr):

    @classmethod
    def validate_half_and_full_email(cls, value: str) -> str:
        if not re.search(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$", value):
            raise ValueError("ERROR.EMAIL_VALIDATION")
        return value

    @classmethod
    def __get_pydantic_core_schema__(cls, source_type: Type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
        email_schema = handler.generate_schema(EmailStr)
        return core_schema.chain_schema(
            [
                email_schema,
                core_schema.no_info_after_validator_function(
                    cls.validate_half_and_full_email,
                    core_schema.str_schema()
                ),
            ]
        )

このように定義しておくことで、各クラスで別々に追加バリデーションを定義する必要がなくなります。もちろん、処理を追加している分処理が遅くはなっているので、データが登録されうる最小限の箇所でのみ使用し、それ以外はデフォルトのEmailStrをそのまま使用しても問題ありません。

class _Test(BaseModel):
    email: CustomEmailStr

class TestIsValid:
    @pytest.mark.parametrize(
        "value, expected",
        [
            ("aiueo@example.com", True),
            ("AIUEO@example.com", True),
            ("aiueo+0@example.com", False),
            ("aiueo+あいうえお@example.com", False),
            ("あいうえお+あいうえお@example.com", False),
        ],
    )
    def test_is_valid(self, value, expected):
        if expected:
            _Test(email=value)
        else:
            with pytest.raises(ValidationError):
                _Test(email=value)

ソースコード

終わりに

単純にValidationをかけるパターンはたくさん出てきますが、特定のクラスを継承するパターンして横展開の工数を減らしたい目的の記事は見つかりませんでした。ぜひ活用してください。

参考情報

Discussion