from datetime import timezone

from sqlalchemy import DateTime, TypeDecorator


class AwareDateTime(TypeDecorator):
    impl = DateTime(timezone=True)

    cache_ok = True

    def process_bind_param(self, value, dialect):
        # Ensure the value is aware. If it's naive, assume UTC.
        if value is not None and value.tzinfo is None:
            value = value.replace(tzinfo=timezone.utc)
        return value

    def process_result_value(self, value, dialect):
        # Ensure the value is aware. If it's naive, assume UTC.
        if value is not None and value.tzinfo is None:
            value = value.replace(tzinfo=timezone.utc)
        return value