From 368b4ba0c175a5ac7103659cfc1c33632c67d624 Mon Sep 17 00:00:00 2001 From: irl Date: Fri, 6 Dec 2024 18:08:09 +0000 Subject: [PATCH] lint: add types for custom type to handle existing naive datetimes --- app/models/types.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/app/models/types.py b/app/models/types.py index 7bb978d..b463d52 100644 --- a/app/models/types.py +++ b/app/models/types.py @@ -1,21 +1,24 @@ -from datetime import timezone +from datetime import datetime, timezone +from typing import Optional -from sqlalchemy import DateTime, TypeDecorator +from sqlalchemy.engine import Dialect +from sqlalchemy.types import TypeDecorator, DateTime -class AwareDateTime(TypeDecorator): +class AwareDateTime(TypeDecorator[datetime]): 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. + def process_bind_param( + self, value: Optional[datetime], dialect: Dialect + ) -> Optional[datetime]: 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. + def process_result_value( + self, value: Optional[datetime], dialect: Dialect + ) -> Optional[datetime]: if value is not None and value.tzinfo is None: value = value.replace(tzinfo=timezone.utc) return value