pydantic_utcdatetime.py

"""Configure pydantic to allow UTCDateTime attributes on models.
"""
from datetime import datetime
from typing import Any, Callable, Dict, List, Tuple, TypeVar, Union

from obspy import UTCDateTime
from pydantic.errors import PydanticValueError
import pydantic.json
import pydantic.schema
import pydantic.validators


# placeholder type for register_custom_pydantic_type method
CustomType = TypeVar("CustomType")


def register_custom_pydantic_type(
    custom_type: CustomType,
    encoder: Callable[[CustomType], Any],
    json_schema: Dict,
    parsers: List[Callable[[Any], CustomType]],
):
    try:
        if custom_type.__custom_pydantic_type__:
            # already registered
            return
    except AttributeError:
        # not registered yet
        pass
    # add encoder
    pydantic.json.ENCODERS_BY_TYPE[custom_type] = encoder
    # add openapi mapping
    pydantic.schema.field_class_to_schema += ((custom_type, json_schema),)
    # add validator
    pydantic.validators._VALIDATORS.append((custom_type, parsers))
    # mark as installed
    custom_type.__custom_pydantic_type__ = True


class UTCDateTimeError(PydanticValueError):
    msg_template = "invalid date-time format"


def format_utcdatetime(o: UTCDateTime) -> str:
    return o.strftime("%Y-%m-%dT%H:%M:%S.%fZ")


def parse_utcdatetime(
    value: Union[datetime, float, int, str, UTCDateTime, Dict]
) -> UTCDateTime:
    if isinstance(value, UTCDateTime):
        # If the value is already a UTCDateTime, return it directly
        return value
    elif isinstance(value, dict):
        # Handle the nanosecond dictionary format
        try:
            # Assuming the dictionary format is something like {'_UTCDateTime__ns': 1598918401000000000}
            ns = value.get("_UTCDateTime__ns")
            if ns is not None:
                return UTCDateTime(ns / 1e9)  # Convert nanoseconds to seconds
        except Exception as e:
            raise UTCDateTimeError() from e
    else:
        # Handle other formats (datetime, float, int, str)
        try:
            return UTCDateTime(value)
        except Exception:
            raise UTCDateTimeError()

    raise UTCDateTimeError()  # Raise an error if none of the above conditions are met


# Code below is what was used to parse UTCDateTime value for pydantic before version 109 break.
# def parse_utcdatetime(
#     value: Union[datetime, float, int, str, UTCDateTime]
# ) -> UTCDateTime:
#     try:
#         return UTCDateTime(value)
#     except:
#         raise UTCDateTimeError()


register_custom_pydantic_type(
    UTCDateTime,
    encoder=format_utcdatetime,
    json_schema={"type": "string", "format": "date-time"},
    parsers=[parse_utcdatetime],
)