Re-infer statements that got more specific on subsequent pass

Fixed issue where mypy plugin would not correctly interpret an explicit
:class:`_orm.Mapped` annotation in conjunction with a
:func:`_orm.relationship` that refers to a class by string name; the
correct annotation would be downgraded to a less specific one leading to
typing errors.

The thing figured out here is that after we've already scanned
a class in the semanal stage and created DeclClassApplied,
when we are called again with that same DeclClassApplied, for this
specific kind of case we actually now have *better* types than
we did before, where the left side that looked like
List?[Address?] now seems to say
builtins.list[official.module.Address] - so let's take the
right side expression again, this time embedded in our
Mapped._empty_constructor() expression, and run the infer
all over again just like mypy would.   Just not setting the
"wrong" type here fixed the test cases but by re-applying the
whole infer we get the correct Mapped[] on the left side too.

Fixes: #6255
Change-Id: Iafe7254374f685a8458c7a1db82aafc2ed6d0232

1 files changed, 46 insertions, 2 deletions

diff --git a/lib/sqlalchemy/ext/mypy/apply.py b/lib/sqlalchemy/ext/mypy/apply.py
index 366260437..5dc9ec0b1 100644
--- a/lib/sqlalchemy/ext/mypy/apply.py
+++ b/lib/sqlalchemy/ext/mypy/apply.py
@@ -15,6 +15,7 @@ from mypy.nodes import AssignmentStmt
 from mypy.nodes import CallExpr
 from mypy.nodes import ClassDef
 from mypy.nodes import MDEF
+from mypy.nodes import MemberExpr
 from mypy.nodes import NameExpr
 from mypy.nodes import StrExpr
 from mypy.nodes import SymbolTableNode
@@ -32,6 +33,7 @@ from mypy.types import TypeOfAny
 from mypy.types import UnboundType
 from mypy.types import UnionType
 
+from . import infer
 from . import util
 
 
@@ -92,6 +94,7 @@ def _re_apply_declarative_assignments(
     mapped_attr_lookup = {
         name: typ for name, typ in cls_metadata.mapped_attr_names
     }
+    update_cls_metadata = False
 
     for stmt in cls.defs.body:
         # for a re-apply, all of our statements are AssignmentStmt;
@@ -104,10 +107,51 @@ def _re_apply_declarative_assignments(
             and stmt.lvalues[0].name in mapped_attr_lookup
             and isinstance(stmt.lvalues[0].node, Var)
         ):
-            typ = mapped_attr_lookup[stmt.lvalues[0].name]
+
             left_node = stmt.lvalues[0].node
+            python_type_for_type = mapped_attr_lookup[stmt.lvalues[0].name]
+            # if we have scanned an UnboundType and now there's a more
+            # specific type than UnboundType, call the re-scan so we
+            # can get that set up correctly
+            if (
+                isinstance(python_type_for_type, UnboundType)
+                and not isinstance(left_node.type, UnboundType)
+                and (
+                    isinstance(stmt.rvalue.callee, MemberExpr)
+                    and stmt.rvalue.callee.expr.node.fullname
+                    == "sqlalchemy.orm.attributes.Mapped"
+                    and stmt.rvalue.callee.name == "_empty_constructor"
+                    and isinstance(stmt.rvalue.args[0], CallExpr)
+                )
+            ):
+
+                python_type_for_type = (
+                    infer._infer_type_from_right_hand_nameexpr(
+                        api,
+                        stmt,
+                        left_node,
+                        left_node.type,
+                        stmt.rvalue.args[0].callee,
+                    )
+                )
+
+                if python_type_for_type is None or isinstance(
+                    python_type_for_type, UnboundType
+                ):
+                    continue
+
+                # update the DeclClassApplied with the better information
+                mapped_attr_lookup[stmt.lvalues[0].name] = python_type_for_type
+                update_cls_metadata = True
+
+            left_node.type = api.named_type(
+                "__sa_Mapped", [python_type_for_type]
+            )
 
-            left_node.type = api.named_type("__sa_Mapped", [typ])
+    if update_cls_metadata:
+        cls_metadata.mapped_attr_names[:] = [
+            (k, v) for k, v in mapped_attr_lookup.items()
+        ]
 
 
 def _apply_type_to_mapped_statement(