From bab4f6555e15b43114ffaadf8d8b98851130795d Mon Sep 17 00:00:00 2001 From: DYNOSuprovo Date: Tue, 29 Sep 2026 16:52:04 +0530 Subject: [PATCH] Use AwareDateTime when generating field from DateTime(timezone=True) (#407) Signed-off-by: DYNOSuprovo --- src/marshmallow_sqlalchemy/convert.py | 14 +++++++++++++- tests/test_conversion.py | 22 +++++++++++++++++++++- tests/test_sqlalchemy_schema.py | 17 +++++++++++++++++ 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/src/marshmallow_sqlalchemy/convert.py b/src/marshmallow_sqlalchemy/convert.py index 2b7c670..656c425 100644 --- a/src/marshmallow_sqlalchemy/convert.py +++ b/src/marshmallow_sqlalchemy/convert.py @@ -1,5 +1,6 @@ from __future__ import annotations +import datetime as dt import functools import inspect import uuid @@ -93,6 +94,16 @@ def _enum_field_factory( ) +def _datetime_field_factory( + converter: ModelConverter, data_type: sa.DateTime +) -> type[fields.Field]: + if getattr(data_type, "timezone", False) or isinstance( + data_type, mssql.DATETIMEOFFSET + ): + return fields.AwareDateTime + return converter.type_mapping.get(dt.datetime, fields.DateTime) + + class ModelConverter: """Converts a SQLAlchemy model into a dictionary of corresponding marshmallow `Fields `. @@ -105,6 +116,7 @@ class ModelConverter: sa.JSON: fields.Raw, sa.ARRAY: _list_field_factory, sa.PickleType: fields.Raw, + sa.DateTime: _datetime_field_factory, postgresql.BIT: fields.Integer, postgresql.OID: fields.Integer, postgresql.UUID: fields.UUID, @@ -123,7 +135,7 @@ class ModelConverter: mysql.SET: fields.List, mysql.ENUM: fields.Field, mysql.INTEGER: fields.Integer, - mysql.DATETIME: fields.DateTime, + mysql.DATETIME: _datetime_field_factory, mssql.BIT: fields.Integer, mssql.UNIQUEIDENTIFIER: fields.UUID, } diff --git a/tests/test_conversion.py b/tests/test_conversion.py index dd467cc..0565ba1 100644 --- a/tests/test_conversion.py +++ b/tests/test_conversion.py @@ -7,7 +7,7 @@ import sqlalchemy as sa from marshmallow import Schema, fields, validate from sqlalchemy import Integer, String -from sqlalchemy.dialects import mysql, postgresql +from sqlalchemy.dialects import mssql, mysql, postgresql from sqlalchemy.orm import Mapped, Session, column_property from marshmallow_sqlalchemy import ( @@ -205,6 +205,26 @@ def test_convert_types(self, converter, sa_type, field_type): field = converter.property2field(prop) assert type(field) is field_type + @pytest.mark.parametrize( + ("sa_type_instance", "field_type"), + ( + (sa.DateTime(timezone=True), fields.AwareDateTime), + (sa.DateTime(timezone=False), fields.DateTime), + (sa.DateTime(), fields.DateTime), + (postgresql.TIMESTAMP(timezone=True), fields.AwareDateTime), + (postgresql.TIMESTAMP(timezone=False), fields.DateTime), + (mysql.DATETIME(timezone=True), fields.AwareDateTime), + (mysql.DATETIME(timezone=False), fields.DateTime), + (mssql.DATETIMEOFFSET(), fields.AwareDateTime), + (sa.TIMESTAMP(timezone=True), fields.AwareDateTime), + (sa.TIMESTAMP(timezone=False), fields.DateTime), + ), + ) + def test_convert_datetime_types(self, converter, sa_type_instance, field_type): + prop = make_property(sa_type_instance) + field = converter.property2field(prop) + assert type(field) is field_type + def test_convert_Numeric(self, converter): prop = make_property(sa.Numeric(scale=2)) field = converter.property2field(prop) diff --git a/tests/test_sqlalchemy_schema.py b/tests/test_sqlalchemy_schema.py index 386e89c..ce3ff8a 100644 --- a/tests/test_sqlalchemy_schema.py +++ b/tests/test_sqlalchemy_schema.py @@ -750,3 +750,20 @@ class Meta: assert schema.dump(teacher) == { "data": {"foo": "bar"}, } + + +def test_auto_field_aware_datetime(Base): + class Item(Base): + __tablename__ = "item_aware_datetime" + id = sa.Column(sa.Integer, primary_key=True) + created_at = sa.Column(sa.DateTime(timezone=True), nullable=False) + updated_at = sa.Column(sa.DateTime(timezone=False), nullable=True) + + class ItemSchema(SQLAlchemyAutoSchema): + class Meta: + model = Item + + schema = ItemSchema() + assert isinstance(schema.fields["created_at"], fields.AwareDateTime) + assert isinstance(schema.fields["updated_at"], fields.DateTime) + assert not isinstance(schema.fields["updated_at"], fields.AwareDateTime)