fixes for mypy, error in stubsfix_mypy

There's a critical mistake in the stubs that we may want to fix.
Or we might just keep it that way as this is all obsolete
anyway.  not sure yet.

Change-Id: I380a1dc177de571b94a176e0f5e5b5d8d4ea2453

4 files changed, 23 insertions, 1 deletions

diff --git a/doc/build/changelog/unreleased_14/mypy_fix.rst b/doc/build/changelog/unreleased_14/mypy_fix.rst
new file mode 100644
index 000000000..9f6d1bcdd
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/mypy_fix.rst
@@ -0,0 +1,6 @@
+.. change::
+    :tags: bug, mypy
+
+    Fixed some issues in the mypy plugin that would otherwise fail
+    with some fixes that are occurring in sqlalchemy2-stubs concurrently.
+
diff --git a/lib/sqlalchemy/ext/mypy/apply.py b/lib/sqlalchemy/ext/mypy/apply.py
index 99be194cd..745a8d759 100644
--- a/lib/sqlalchemy/ext/mypy/apply.py
+++ b/lib/sqlalchemy/ext/mypy/apply.py
@@ -44,6 +44,7 @@ def apply_mypy_mapped_attr(
     api: SemanticAnalyzerPluginInterface,
     item: Union[NameExpr, StrExpr],
     attributes: List[util.SQLAlchemyAttribute],
+    reset_statement_type: bool = False,
 ) -> None:
     if isinstance(item, NameExpr):
         name = item.name
@@ -91,6 +92,11 @@ def apply_mypy_mapped_attr(
         api, stmt, stmt.lvalues[0], left_hand_explicit_type, None
     )
 
+    if reset_statement_type:
+        stmt.type = api.named_type(
+            NAMED_TYPE_SQLA_MAPPED, [AnyType(TypeOfAny.special_form)]
+        )
+
 
 def re_apply_declarative_assignments(
     cls: ClassDef,
@@ -226,6 +232,10 @@ def apply_type_to_mapped_statement(
     # internally
     stmt.rvalue = util.expr_to_mapped_constructor(stmt.rvalue)
 
+    # this works for every test except the dataclasses one
+    if stmt.type is not None and python_type_for_type is not None:
+        stmt.type = python_type_for_type
+
 
 def add_additional_orm_attributes(
     cls: ClassDef,
diff --git a/lib/sqlalchemy/ext/mypy/decl_class.py b/lib/sqlalchemy/ext/mypy/decl_class.py
index c33c30e25..3c9847b73 100644
--- a/lib/sqlalchemy/ext/mypy/decl_class.py
+++ b/lib/sqlalchemy/ext/mypy/decl_class.py
@@ -398,7 +398,9 @@ def _scan_declarative_assignment_stmt(
         else:
             for item in stmt.rvalue.items:
                 if isinstance(item, (NameExpr, StrExpr)):
-                    apply.apply_mypy_mapped_attr(cls, api, item, attributes)
+                    apply.apply_mypy_mapped_attr(
+                        cls, api, item, attributes, reset_statement_type=True
+                    )
 
     left_hand_mapped_type: Optional[Type] = None
     left_hand_explicit_type: Optional[ProperType] = None
diff --git a/test/ext/mypy/test_mypy_plugin_py3k.py b/test/ext/mypy/test_mypy_plugin_py3k.py
index 3df758c56..cb04d1c73 100644
--- a/test/ext/mypy/test_mypy_plugin_py3k.py
+++ b/test/ext/mypy/test_mypy_plugin_py3k.py
@@ -76,6 +76,10 @@ class MypyPluginTest(fixtures.TestBase):
                 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