Custom UTC DateTime Type

[sf-wilbur]

I am trying to create a script that will allow me to use UTCDateTime when processing geomagnetic data as I was doing so with pydantic version 1 and the json_encoders function. I updated to Pydantic version 2 and changed the necessary functions and method names thorughout my whole program but am unable to use the new tools, such as annotated or generate_core_schema to allow the Pydantic Basemodel to accept UTCDatetimes. I need to be able to provide UTCDatetime to multiple Pydantic basemodels throughout my program and would prefer to not have to iteratively go through my whole program and change each instance of UTCDatetime to a new defined custom type.

Here is what I have tried but this only generates an unable to generate core_scheme for obspy.utcdatatime.utcdatetime error.

`from datetime import datetime
from typing import Any, List
from obspy.core.utcdatetime import UTCDateTime
from pydantic import BaseModel, GetCoreSchemaHandler
from pydantic_core import CoreSchema, core_schema

class CustomDate(BaseModel):
"""Custom date handling for strings and ObsPy's UTCDateTime."""
allowed_formats: List[str] = ['%Y-%m-%d', '%d-%m-%Y', '%m/%d/%Y']

class config:
    arbitrary_types_allowed = True 
    json_encoders = {
    UTCDateTime: lambda v: v.isoformat()  # Serializing UTCDateTime to ISO format string
}

@classmethod
def try_parse_date(cls, v: Any, info: core_schema.ValidationInfo) -> Any:
    if isinstance(v, str):
        for fmt in cls.allowed_formats:
            try:
                return cls(datetime.strptime(v, fmt).date())
            except ValueError:
                continue
        raise ValueError("Unable to parse date from string using allowed formats.")
    
    if isinstance(v, UTCDateTime):
        return cls(v.datetime.date())
    else:
        raise ValueError("Unsupported type for date conversion.")

@classmethod
def __get_pydantic_core_schema__(
    cls, source_type: Any, handler: GetCoreSchemaHandler) ->CoreSchema:
    
    # assert source_type is UTCDateTime
    return core_schema.with_info_after_validator_function(
    cls.model_validate, handler.generate_schema(UTCDateTime),
    core_schema.str_schema(), 
    serialization=core_schema.to_string_ser_schema()
    )

Example usage

if name == "main":
custom_date = CustomDate.try_parse_date("2023-03-21")
print(custom_date)

`

Here is my original code that was working before updating to version 2 I added the version 1 handles to the Pydantic tools as an interim fix but that isn't working either.

`"""Configure pydantic to allow UTCDateTime attributes on models.
"""

from datetime import datetime
from typing import Any, Callable, Dict, List, Tuple, TypeVar, Union

FIXME: migrating to pydantic v2 should make this unnecessary

(other changes will be needed though)

from fastapi.encoders import ENCODERS_BY_TYPE

from obspy import UTCDateTime
from pydantic.v1.errors import PydanticValueError
import pydantic.v1.json
import pydantic.v1.schema
import pydantic.v1.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.v1.json.ENCODERS_BY_TYPE[custom_type] = encoder

# FIXME: migrating to pydantic v2 should make this unnecessary
#        (other changes will be needed though)
ENCODERS_BY_TYPE[custom_type] = encoder

# add openapi mapping
pydantic.v1.schema.field_class_to_schema += ((custom_type, json_schema),)
# add validator
pydantic.v1.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]
) -> 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],
)
`