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

diff --git a/doc/build/changelog/unreleased_14/6937.rst b/doc/build/changelog/unreleased_14/6937.rst
new file mode 100644 (file)
index 0000000..5dfadd5
--- /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 23c78aa51fc47b12c063d25a0eeda875fe4cf5ee..b85ec0f699e10c6081f6e87f83a45645e2d54c55 100644 (file)
--- 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 614805d77c139fbea2ca116c5f4c207ffb24f1d3..a3825f175f65d2bb2391b8f590bbb122e6d1e49b 100644 (file)
--- 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 (file)
index 0000000..cfbe83d
--- /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)