mypy plugin fixes

Adjustments made to the mypy plugin to accommodate for some potential
changes being made for issue #236 sqlalchemy2-stubs when using SQLAlchemy
1.4. These changes are being kept in sync within SQLAlchemy 2.0.
The changes are also backwards compatible with older versions of
sqlalchemy2-stubs.

Fixed crash in mypy plugin which could occur on both 1.4 and 2.0 versions
if a decorator for the :func:`_orm.registry.mapped` decorator were used
that was referenced in an expression with more than two components (e.g.
``@Backend.mapper_registry.mapped``). This scenario is now ignored; when
using the plugin, the decorator expression needs to be two components (i.e.
``@reg.mapped``).

References: sqlalchemy/sqlalchemy2-stubs#236
Fixes: #9102
Change-Id: Ieb1bf7bf8184645bcd43253e57f1c267b2640537
(cherry picked from commit cf64582f61b15716228302f669322d7efa1003c1)
(cherry picked from commit 36285760238314f70eed4532ca2c2c0c2d684609)

doc/build/changelog/unreleased_14/mypy_fix.rst

@@ -0,0 +1,22 @@
+.. change::
+    :tags: bug, mypy
+    :versions: 2.0.0rc3
+
+    Adjustments made to the mypy plugin to accommodate for some potential
+    changes being made for issue #236 sqlalchemy2-stubs when using SQLAlchemy
+    1.4. These changes are being kept in sync within SQLAlchemy 2.0.
+    The changes are also backwards compatible with older versions of
+    sqlalchemy2-stubs.
+
+
+.. change::
+    :tags: bug, mypy
+    :tickets: 9102
+    :versions: 2.0.0rc3
+
+    Fixed crash in mypy plugin which could occur on both 1.4 and 2.0 versions
+    if a decorator for the :func:`_orm.registry.mapped` decorator were used
+    that was referenced in an expression with more than two components (e.g.
+    ``@Backend.mapper_registry.mapped``). This scenario is now ignored; when
+    using the plugin, the decorator expression needs to be two components (i.e.
+    ``@reg.mapped``).

lib/sqlalchemy/ext/mypy/apply.py

@@ -164,7 +164,10 @@ def re_apply_declarative_assignments(
 
                 update_cls_metadata = True
 
-            if python_type_for_type is not None:
+            if python_type_for_type is not None and (
+                not isinstance(left_node.type, Instance)
+                or left_node.type.type.fullname != NAMED_TYPE_SQLA_MAPPED
+            ):
                 left_node.type = api.named_type(
                     NAMED_TYPE_SQLA_MAPPED, [python_type_for_type]
                 )
@@ -201,15 +204,23 @@ class User(Base):
     left_node = lvalue.node
     assert isinstance(left_node, Var)
 
+    # to be completely honest I have no idea what the difference between
+    # left_node.type and stmt.type is, what it means if these are different
+    # vs. the same, why in order to get tests to pass I have to assign
+    # to stmt.type for the second case and not the first.  this is complete
+    # trying every combination until it works stuff.
+
     if left_hand_explicit_type is not None:
         left_node.type = api.named_type(
             NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
         )
     else:
         lvalue.is_inferred_def = False
-        left_node.type = api.named_type(
+        left_node.type = stmt.type = api.named_type(
             NAMED_TYPE_SQLA_MAPPED,
-            [] if python_type_for_type is None else [python_type_for_type],
+            [AnyType(TypeOfAny.special_form)]
+            if python_type_for_type is None
+            else [python_type_for_type],
         )
 
     # so to have it skip the right side totally, we can do this:
@@ -226,6 +237,11 @@ class User(Base):
     # internally
     stmt.rvalue = util.expr_to_mapped_constructor(stmt.rvalue)
 
+    if stmt.type is None or python_type_for_type is None:
+        stmt.type = api.named_type(
+            NAMED_TYPE_SQLA_MAPPED, [AnyType(TypeOfAny.special_form)]
+        )
+
 
 def add_additional_orm_attributes(
     cls: ClassDef,

lib/sqlalchemy/ext/mypy/plugin.py

@@ -184,10 +184,13 @@ def _fill_in_decorators(ctx: ClassDefContext) -> None:
         else:
             continue
 
-        assert isinstance(target.expr, NameExpr)
-        sym = ctx.api.lookup_qualified(
-            target.expr.name, target, suppress_errors=True
-        )
+        if isinstance(target.expr, NameExpr):
+            sym = ctx.api.lookup_qualified(
+                target.expr.name, target, suppress_errors=True
+            )
+        else:
+            continue
+
         if sym and sym.node:
             sym_type = get_proper_type(sym.type)
             if isinstance(sym_type, Instance):

test/ext/mypy/files/issue_9102.py

@@ -0,0 +1,18 @@
+from sqlalchemy import Column
+from sqlalchemy import Integer
+from sqlalchemy.orm import registry
+
+
+class BackendMeta:
+    __abstract__ = True
+    mapped_registry: registry = registry()
+    metadata = mapped_registry.metadata
+
+
+# this decorator is not picked up now, but at least it doesn't crash
+@BackendMeta.mapped_registry.mapped
+class User:
+    __tablename__ = "user"
+
+    # EXPECTED_MYPY: Incompatible types in assignment (expression has type "Column[Integer]", variable has type "int")  # noqa: E501
+    id: int = Column(Integer(), primary_key=True)

test/ext/mypy/files/issue_9102_workaround.py

@@ -0,0 +1,19 @@
+from sqlalchemy import Column
+from sqlalchemy import Integer
+from sqlalchemy.orm import registry
+
+
+class BackendMeta:
+    __abstract__ = True
+    mapped_registry: registry = registry()
+    metadata = mapped_registry.metadata
+
+
+reg: registry = BackendMeta.mapped_registry
+
+
+@reg.mapped
+class User:
+    __tablename__ = "user"
+
+    id: int = Column(Integer(), primary_key=True)

test/ext/mypy/test_mypy_plugin_py3k.py

@@ -76,6 +76,10 @@ def run(path, use_plugin=True, incremental=False):
                 shutil.copyfile(path, test_program)
                 args.append(test_program)
 
+            # I set this locally but for the suite here needs to be
+            # disabled
+            os.environ.pop("MYPY_FORCE_COLOR", None)
+
             result = api.run(args)
             return result