lint: add types for custom type to handle existing naive datetimes

This commit is contained in:
Iain Learmonth 2024-12-06 18:08:09 +00:00
parent 39bdac1ecf
commit 368b4ba0c1

View file

@ -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) impl = DateTime(timezone=True)
cache_ok = True cache_ok = True
def process_bind_param(self, value, dialect): def process_bind_param(
# Ensure the value is aware. If it's naive, assume UTC. self, value: Optional[datetime], dialect: Dialect
) -> Optional[datetime]:
if value is not None and value.tzinfo is None: if value is not None and value.tzinfo is None:
value = value.replace(tzinfo=timezone.utc) value = value.replace(tzinfo=timezone.utc)
return value return value
def process_result_value(self, value, dialect): def process_result_value(
# Ensure the value is aware. If it's naive, assume UTC. self, value: Optional[datetime], dialect: Dialect
) -> Optional[datetime]:
if value is not None and value.tzinfo is None: if value is not None and value.tzinfo is None:
value = value.replace(tzinfo=timezone.utc) value = value.replace(tzinfo=timezone.utc)
return value return value