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

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