Add test for declared_attr

pamelafox · pamelafox · commit cf4170adcdae · 2023-08-25T16:28:51.000-07:00
diff --git a/docs/customizing.rst b/docs/customizing.rst
@@ -21,29 +21,26 @@ joined-table inheritance.
 
 .. code-block:: python
 
-    from flask_sqlalchemy.model import Model
-    import sqlalchemy as sa
-    import sqlalchemy.orm
+    from sqlalchemy import Integer, String, ForeignKey
+    from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, declared_attr
 
-    class IdModel(Model):
-        @sa.orm.declared_attr
+    class Base(DeclarativeBase):
+        @declared_attr.cascading
+        @classmethod
         def id(cls):
             for base in cls.__mro__[1:-1]:
                 if getattr(base, "__table__", None) is not None:
-                    type = sa.ForeignKey(base.id)
-                    break
-            else:
-                type = sa.Integer
+                        return mapped_column(ForeignKey(base.id), primary_key=True)
+                else:
+                    return mapped_column(Integer, primary_key=True)
 
-            return sa.Column(type, primary_key=True)
-
-    db = SQLAlchemy(model_class=IdModel)
+    db = SQLAlchemy(app, model_class=Base)
 
     class User(db.Model):
-        name = db.Column(db.String)
+        name: Mapped[str] = mapped_column(String)
 
     class Employee(User):
-        title = db.Column(db.String)
+        title: Mapped[str] = mapped_column(String)
 
 
 Abstract Models and Mixins
@@ -56,31 +53,49 @@ they are created or updated.
 .. code-block:: python
 
     from datetime import datetime
+    from sqlalchemy import DateTime, Integer, String
+    from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, declared_attr
 
     class TimestampModel(db.Model):
         __abstract__ = True
-        created: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, default=datetime.utcnow)
-        updated: Mapped[datetime] = mapped_column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+        created: Mapped[datetime] = mapped_column(DateTime, nullable=False, default=datetime.utcnow)
+        updated: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
 
     class Author(db.Model):
-        id: Mapped[int] = mapped_column(db.Integer, primary_key=True)
-        username: Mapped[str] = mapped_column(db.String, unique=True, nullable=False)
+        id: Mapped[int] = mapped_column(Integer, primary_key=True)
+        username: Mapped[str] = mapped_column(String, unique=True, nullable=False)
 
     class Post(TimestampModel):
-        id: Mapped[int] = mapped_column(db.Integer, primary_key=True)
-        title: Mapped[str] = mapped_column(db.String, nullable=False)
+        id: Mapped[int] = mapped_column(Integer, primary_key=True)
+        title: Mapped[str] = mapped_column(String, nullable=False)
 
 This can also be done with a mixin class, inheriting from ``db.Model`` separately.
 
 .. code-block:: python
 
-    class TimestampModel:
-        created: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, default=datetime.utcnow)
-        updated: Mapped[datetime] = mapped_column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+    class TimestampMixin:
+        created: Mapped[datetime] = mapped_column(DateTime, nullable=False, default=datetime.utcnow)
+        updated: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+
+    class Post(TimestampMixin, db.Model):
+        id: Mapped[int] = mapped_column(Integer, primary_key=True)
+        title: Mapped[str] = mapped_column(String, nullable=False)
+
 
-    class Post2(TimestampModel, db.Model):
-        id: Mapped[int] = mapped_column(db.Integer, primary_key=True)
-        title: Mapped[str] = mapped_column(db.String, nullable=False)
+Disabling Table Name Generation
+-------------------------------
+
+Some projects prefer to set each model's ``__tablename__`` manually rather than relying
+on Flask-SQLAlchemy's detection and generation. The simple way to achieve that is to
+set each ``__tablename__`` and not modify the base class. However, the table name
+generation can be disabled by setting `disable_autonaming=True` in the `SQLAlchemy` constructor.
+
+.. code-block:: python
+
+    class Base(sa_orm.DeclarativeBase):
+        pass
+
+    db = SQLAlchemy(app, model_class=Base, disable_autonaming=True)
 
 
 Session Class
@@ -161,19 +176,3 @@ To customize only ``session.query``, pass the ``query_cls`` key to the
 .. code-block:: python
 
     db = SQLAlchemy(session_options={"query_cls": GetOrQuery})
-
-
-Disabling Table Name Generation
--------------------------------
-
-Some projects prefer to set each model's ``__tablename__`` manually rather than relying
-on Flask-SQLAlchemy's detection and generation. The simple way to achieve that is to
-set each ``__tablename__`` and not modify the base class. However, the table name
-generation can be disabled by setting `disable_autonaming=True` in the `SQLAlchemy` constructor.
-
-.. code-block:: python
-
-    class Base(sa_orm.DeclarativeBase):
-        pass
-
-    db = SQLAlchemy(app, model_class=Base, disable_autonaming=True)
diff --git a/src/flask_sqlalchemy/extension.py b/src/flask_sqlalchemy/extension.py
@@ -540,10 +540,8 @@ def _make_declarative_base(
                 " Got: {}".format(model_class.__bases__)
             )
         elif len(declarative_bases) == 1:
-            body = {
-                "__fsa__": self,
-                "metadata": model_class.metadata,  # type: ignore[attr-defined]
-            }
+            body = dict(model_class.__dict__)  # type: ignore[arg-type]
+            body["__fsa__"] = self
             mixin_classes = [BindMixin, NameMixin, Model]
             if disable_autonaming:
                 mixin_classes.remove(NameMixin)
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+from datetime import datetime
 
 import pytest
 import sqlalchemy as sa
@@ -80,6 +81,182 @@ class Base(sa_orm.DeclarativeBaseNoMeta, sa_orm.MappedAsDataclass):
     assert isinstance(db.Model, sa_orm.decl_api.DCTransformDeclarative)
 
 
+@pytest.mark.usefixtures("app_ctx")
+def test_declaredattr(app: Flask, model_class: t.Any) -> None:
+    if model_class is Model:
+
+        class IdModel(Model):
+            @sa.orm.declared_attr
+            @classmethod
+            def id(cls: type[Model]):  # type: ignore[no-untyped-def]
+                for base in cls.__mro__[1:-1]:
+                    if getattr(base, "__table__", None) is not None and hasattr(
+                        base, "id"
+                    ):
+                        return sa.Column(sa.ForeignKey(base.id), primary_key=True)
+                return sa.Column(sa.Integer, primary_key=True)
+
+        db = SQLAlchemy(app, model_class=IdModel)
+
+        class User(db.Model):
+            name = db.Column(db.String)
+
+        class Employee(User):
+            title = db.Column(db.String)
+
+    else:
+
+        class Base(sa_orm.DeclarativeBase):
+            @sa_orm.declared_attr
+            @classmethod
+            def id(cls: type[sa_orm.DeclarativeBase]) -> sa_orm.Mapped[int]:
+                for base in cls.__mro__[1:-1]:
+                    if getattr(base, "__table__", None) is not None and hasattr(
+                        base, "id"
+                    ):
+                        return sa_orm.mapped_column(
+                            db.ForeignKey(base.id), primary_key=True
+                        )
+                return sa_orm.mapped_column(db.Integer, primary_key=True)
+
+        db = SQLAlchemy(app, model_class=Base)
+
+        class User(db.Model):  # type: ignore[no-redef]
+            name: sa_orm.Mapped[str] = sa_orm.mapped_column(db.String)
+
+        class Employee(User):  # type: ignore[no-redef]
+            title: sa_orm.Mapped[str] = sa_orm.mapped_column(db.String)
+
+    db.create_all()
+    db.session.add(Employee(name="Emp Loyee", title="Admin"))
+    db.session.commit()
+    user = db.session.execute(db.select(User)).scalar()
+    employee = db.session.execute(db.select(Employee)).scalar()
+    assert user is not None
+    assert employee is not None
+    assert user.id == 1
+    assert employee.id == 1
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_abstractmodel(app: Flask, model_class: t.Any) -> None:
+    db = SQLAlchemy(app, model_class=model_class)
+
+    if issubclass(db.Model, (sa_orm.MappedAsDataclass)):
+
+        class TimestampModel(db.Model):
+            __abstract__ = True
+            created: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime, nullable=False, insert_default=datetime.utcnow, init=False
+            )
+            updated: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime,
+                insert_default=datetime.utcnow,
+                onupdate=datetime.utcnow,
+                init=False,
+            )
+
+        class Post(TimestampModel):
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(
+                db.Integer, primary_key=True, init=False
+            )
+            title: sa_orm.Mapped[str] = sa_orm.mapped_column(db.String, nullable=False)
+
+    elif issubclass(db.Model, (sa_orm.DeclarativeBase, sa_orm.DeclarativeBaseNoMeta)):
+
+        class TimestampModel(db.Model):  # type: ignore[no-redef]
+            __abstract__ = True
+            created: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime, nullable=False, default=datetime.utcnow
+            )
+            updated: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow
+            )
+
+        class Post(TimestampModel):  # type: ignore[no-redef]
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(db.Integer, primary_key=True)
+            title: sa_orm.Mapped[str] = sa_orm.mapped_column(db.String, nullable=False)
+
+    else:
+
+        class TimestampModel(db.Model):  # type: ignore[no-redef]
+            __abstract__ = True
+            created = db.Column(db.DateTime, nullable=False, default=datetime.utcnow)
+            updated = db.Column(
+                db.DateTime, onupdate=datetime.utcnow, default=datetime.utcnow
+            )
+
+        class Post(TimestampModel):  # type: ignore[no-redef]
+            id = db.Column(db.Integer, primary_key=True)
+            title = db.Column(db.String, nullable=False)
+
+    db.create_all()
+    db.session.add(Post(title="Admin Post"))
+    db.session.commit()
+    post = db.session.execute(db.select(Post)).scalar()
+    assert post is not None
+    assert post.created is not None
+    assert post.updated is not None
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_mixinmodel(app: Flask, model_class: t.Any) -> None:
+    db = SQLAlchemy(app, model_class=model_class)
+
+    if issubclass(db.Model, (sa_orm.MappedAsDataclass)):
+
+        class TimestampMixin(sa_orm.MappedAsDataclass):
+            created: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime, nullable=False, insert_default=datetime.utcnow, init=False
+            )
+            updated: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime,
+                insert_default=datetime.utcnow,
+                onupdate=datetime.utcnow,
+                init=False,
+            )
+
+        class Post(TimestampMixin, db.Model):
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(
+                db.Integer, primary_key=True, init=False
+            )
+            title: sa_orm.Mapped[str] = sa_orm.mapped_column(db.String, nullable=False)
+
+    elif issubclass(db.Model, (sa_orm.DeclarativeBase, sa_orm.DeclarativeBaseNoMeta)):
+
+        class TimestampMixin:  # type: ignore[no-redef]
+            created: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime, nullable=False, default=datetime.utcnow
+            )
+            updated: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+                db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow
+            )
+
+        class Post(TimestampMixin, db.Model):  # type: ignore[no-redef]
+            id: sa_orm.Mapped[int] = sa_orm.mapped_column(db.Integer, primary_key=True)
+            title: sa_orm.Mapped[str] = sa_orm.mapped_column(db.String, nullable=False)
+
+    else:
+
+        class TimestampMixin:  # type: ignore[no-redef]
+            created = db.Column(db.DateTime, nullable=False, default=datetime.utcnow)
+            updated = db.Column(
+                db.DateTime, onupdate=datetime.utcnow, default=datetime.utcnow
+            )
+
+        class Post(TimestampMixin, db.Model):  # type: ignore[no-redef]
+            id = db.Column(db.Integer, primary_key=True)
+            title = db.Column(db.String, nullable=False)
+
+    db.create_all()
+    db.session.add(Post(title="Admin Post"))
+    db.session.commit()
+    post = db.session.execute(db.select(Post)).scalar()
+    assert post is not None
+    assert post.created is not None
+    assert post.updated is not None
+
+
 @pytest.mark.usefixtures("app_ctx")
 def test_model_repr(db: SQLAlchemy) -> None:
     class User(db.Model):
