from typing import Any
from datetime import datetime, timezone, timedelta
from pydantic import BaseModel, AwareDatetime
from pydantic_core import core_schema

# Load GetCoreSchemaHandler via try/except because it was made public in v2.4.0
try:
    from pydantic.annotated_handlers import GetCoreSchemaHandler
except:
    from pydantic._internal._annotated_handlers import GetCoreSchemaHandler


class UTCAwareDatetime(AwareDatetime):
    @classmethod
    def __get_pydantic_core_schema__(
        cls, source: type[Any], handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        if cls is source:
            # used directly as a type
            return core_schema.datetime_schema(tz_constraint=0)
        else:
            schema = handler(source)
            _check_annotated_type(schema["type"], "datetime", cls.__name__)
            schema["tz_constraint"] = 0
            return schema

    def __repr__(self) -> str:
        return "UTCAwareDatetime"


class TestModel(BaseModel):
    created_at: UTCAwareDatetime


now = datetime.now(tz=timezone(timedelta(seconds=1800)))

model = TestModel(created_at=now)
print(model.model_dump_json())

from typing_extensions import Annotated
from datetime import datetime, timezone, timedelta
from pydantic import BaseModel, AwareDatetime
from pydantic.functional_validators import AfterValidator


def check_tz(v: datetime):
    assert v.utcoffset() == timedelta(seconds=0)
    return v


UTCAwareDatetime = Annotated[AwareDatetime, AfterValidator(check_tz)]


class TestModel(BaseModel):
    created_at: UTCAwareDatetime


now = datetime.now(tz=timezone(timedelta(seconds=1800)))

model = TestModel(created_at=now)
print(model.model_dump_json())