Custom Pydantic Types for Dates

custom_pydantic_types.py

import dataclasses
import datetime
import typing

import pendulum
from dataclasses_avroschema.faker import fake
from dataclasses_avroschema.fields import mapper
from dataclasses_avroschema.pydantic.fields import PydanticField
from pydantic import GetJsonSchemaHandler, GetCoreSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema, PydanticCustomError
from typing_extensions import Annotated, Literal


@dataclasses.dataclass
class DateTimeAsString:
    type: Literal["date_as_yyyymmdd", "iso_datetime", "any_valid_date"]

    def __get_pydantic_json_schema(
        self,
        core_schema: core_schema.CoreSchema,  # type: ignore
        handler: GetJsonSchemaHandler
    ) -> JsonSchemaValue:
        field_schema = handler(core_schema)
        return field_schema

    def __get_pydantic_core_schema__(
        self,
        source: typing.Any,
        handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:  # type: ignore
        function_lookup = {
            'date_as_yyyymmdd': typing.cast(core_schema.WithInfoValidatorFunction, self.validate_date),
            'iso_datetime': typing.cast(core_schema.WithInfoValidatorFunction, self.validate_iso_datetime),
            'any_valid_date': typing.cast(core_schema.WithInfoValidatorFunction, self.validate_any_datetime),
        }

        return core_schema.with_info_after_validator_function(
            function_lookup[self.type],
            handler(source),
        )

    @staticmethod
    def validate_any_datetime(datestr: str, _: core_schema.ValidationInfo) -> str:
        try:
            pendulum.parse(datestr, exact=True)
            return datestr
        except ValueError:
            raise PydanticCustomError('invalid_date', 'Not a Valid Date/DateTime')

    @staticmethod
    def validate_iso_datetime(datestr: str, _: core_schema.ValidationInfo) -> str:
        try:
            parsed_date = datetime.datetime.strptime(datestr, "%Y-%m-%dT%H:00")
            return parsed_date.strftime("%Y-%m-%dT%H:00")
        except ValueError:
            raise PydanticCustomError('invalid_date', 'invalid date not in YYYY-MM-DDTHH:00')

    @staticmethod
    def validate_date(datestr: str, _: core_schema.ValidationInfo) -> str:
        try:
            parsed_date = datetime.datetime.strptime(datestr, "%Y-%m-%d")
            return parsed_date.strftime("%Y-%m-%d")
        except ValueError:
            raise PydanticCustomError('invalid_date', 'invalid date not in YYYY-MM-DD')

    def __hash__(self) -> int:
        return hash(type(self.type))


WfDateStr = Annotated[str, DateTimeAsString('date_as_yyyymmdd')]
WfDateTimeStr = Annotated[str, DateTimeAsString('iso_datetime')]
WfAnyValidDateTimeStr = Annotated[str, DateTimeAsString('any_valid_date')]