From 07b7e1594deecc6c9c64893a63f56403da0dcf05 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Wed, 25 Aug 2021 13:12:55 -0400
Subject: [PATCH] ensure "sqlalchemy" info set for all considered classes

Fixed issue in mypy plugin where columns on a mixin would not be correctly
interpreted if the mapped class relied upon a ``__tablename__`` routine
that came from a superclass.

Fixes: #6937
Change-Id: I74aed4862d0545008ee67f781aaa794ab6866926
---
 doc/build/changelog/unreleased_14/6937.rst |  7 ++++++
 lib/sqlalchemy/ext/mypy/decl_class.py      |  3 +++
 lib/sqlalchemy/ext/mypy/util.py            |  4 ++++
 test/ext/mypy/files/mixin_w_tablename.py   | 27 ++++++++++++++++++++++
 4 files changed, 41 insertions(+)
 create mode 100644 doc/build/changelog/unreleased_14/6937.rst
 create mode 100644 test/ext/mypy/files/mixin_w_tablename.py

diff --git a/doc/build/changelog/unreleased_14/6937.rst b/doc/build/changelog/unreleased_14/6937.rst
new file mode 100644
index 0000000000..5dfadd552a
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/6937.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: bug, mypy
+    :tickets: 6937
+
+    Fixed issue in mypy plugin where columns on a mixin would not be correctly
+    interpreted if the mapped class relied upon a ``__tablename__`` routine
+    that came from a superclass.
diff --git a/lib/sqlalchemy/ext/mypy/decl_class.py b/lib/sqlalchemy/ext/mypy/decl_class.py
index 23c78aa51f..b85ec0f699 100644
--- a/lib/sqlalchemy/ext/mypy/decl_class.py
+++ b/lib/sqlalchemy/ext/mypy/decl_class.py
@@ -61,6 +61,9 @@ def scan_declarative_assignments_and_apply_types(
         List[util.SQLAlchemyAttribute]
     ] = util.get_mapped_attributes(info, api)
 
+    # used by assign.add_additional_orm_attributes among others
+    util.establish_as_sqlalchemy(info)
+
     if mapped_attributes is not None:
         # ensure that a class that's mapped is always picked up by
         # its mapped() decorator or declarative metaclass before
diff --git a/lib/sqlalchemy/ext/mypy/util.py b/lib/sqlalchemy/ext/mypy/util.py
index 614805d77c..a3825f175f 100644
--- a/lib/sqlalchemy/ext/mypy/util.py
+++ b/lib/sqlalchemy/ext/mypy/util.py
@@ -99,6 +99,10 @@ def _get_info_mro_metadata(info: TypeInfo, key: str) -> Optional[Any]:
     return None
 
 
+def establish_as_sqlalchemy(info: TypeInfo) -> None:
+    info.metadata.setdefault("sqlalchemy", {})
+
+
 def set_is_base(info: TypeInfo) -> None:
     _set_info_metadata(info, "is_base", True)
 
diff --git a/test/ext/mypy/files/mixin_w_tablename.py b/test/ext/mypy/files/mixin_w_tablename.py
new file mode 100644
index 0000000000..cfbe83d35d
--- /dev/null
+++ b/test/ext/mypy/files/mixin_w_tablename.py
@@ -0,0 +1,27 @@
+# test #6937
+from sqlalchemy import Column
+from sqlalchemy import Integer
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import declared_attr
+from sqlalchemy.orm import Mapped
+
+
+Base = declarative_base()
+
+
+class UpdatedCls:
+    @declared_attr
+    def __tablename__(cls) -> Mapped[str]:
+        return cls.__name__.lower()
+
+    updated_at = Column(Integer)
+
+
+class Bar(UpdatedCls, Base):
+    id = Column(Integer(), primary_key=True)
+    num = Column(Integer)
+
+
+Bar.updated_at.in_([1, 2, 3])
+
+b1 = Bar(num=5, updated_at=6)
-- 
2.47.2