pydantic_utcdatetime.py

"""Add obspy's UTCDateTime to the v2 pydantic schema. 
CustomUTCDateTimeType should be used in place of UTCDateTime on pydantic models.
"""

import datetime
from dateutil import tz
from obspy import UTCDateTime
from pydantic_core import CoreSchema, core_schema
from typing import Annotated, Any

from pydantic import (
    GetCoreSchemaHandler,
    GetJsonSchemaHandler,
)

from pydantic.json_schema import JsonSchemaValue


class CustomUTCDateTimeValidator:
    @classmethod
    def __get_pydantic_core_schema__(
        cls,
        _source_type: Any,
        _handler: GetCoreSchemaHandler,
    ) -> CoreSchema:
        def UTCDateTime_validator(value: Any):
            # if the user inputs an unaware datetime.datetime, make it aware
            if isinstance(value, datetime.datetime):
                if value.tzinfo is not tz.tzutc():
                    value = value.replace(tzinfo=tz.tzutc())
            try:
                time = UTCDateTime(value)
            except:
                raise ValueError(
                    f"Invalid time type. See obspy UTCDateTime for more information."
                )
            return time

        schema = core_schema.chain_schema(
            [
                core_schema.any_schema(),
                core_schema.no_info_plain_validator_function(UTCDateTime_validator),
            ],
        )

        schema = core_schema.json_or_python_schema(
            json_schema=schema,
            python_schema=core_schema.union_schema(
                [
                    # check if it's an instance first before doing any further work
                    core_schema.is_instance_schema(UTCDateTime),
                    schema,
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                lambda instance: instance.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
            ),
        )
        return schema

    @classmethod
    def __get_pydantic_json_schema__(
        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
    ) -> JsonSchemaValue:
        return handler(core_schema.any_schema())


CustomUTCDateTimeType = Annotated[UTCDateTime, CustomUTCDateTimeValidator]