lint: add types for custom type to handle existing naive datetimes
This commit is contained in:
parent
39bdac1ecf
commit
368b4ba0c1
1 changed files with 11 additions and 8 deletions
|
@ -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
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue