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

3 files changed, 103 insertions, 32 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(
diff --git a/lib/sqlalchemy/ext/mypy/decl_class.py b/lib/sqlalchemy/ext/mypy/decl_class.py
index 8fac36342..2870eeb6f 100644
--- a/lib/sqlalchemy/ext/mypy/decl_class.py
+++ b/lib/sqlalchemy/ext/mypy/decl_class.py
@@ -63,10 +63,11 @@ def _scan_declarative_assignments_and_apply_types(
         if not is_mixin_scan:
             assert cls_metadata.is_mapped
 
-            # mypy can call us more than once.  it then will have reset the
+            # mypy can call us more than once.  it then *may* have reset the
             # left hand side of everything, but not the right that we removed,
             # removing our ability to re-scan.   but we have the types
-            # here, so lets re-apply them.
+            # here, so lets re-apply them, or if we have an UnboundType,
+            # we can re-scan
 
             apply._re_apply_declarative_assignments(cls, api, cls_metadata)
 
@@ -422,33 +423,11 @@ def _scan_declarative_assignment_stmt(
         stmt.rvalue.callee, RefExpr
     ):
 
-        type_id = names._type_id_for_callee(stmt.rvalue.callee)
+        python_type_for_type = infer._infer_type_from_right_hand_nameexpr(
+            api, stmt, node, left_hand_explicit_type, stmt.rvalue.callee
+        )
 
-        if type_id is None:
-            return
-        elif type_id is names.COLUMN:
-            python_type_for_type = infer._infer_type_from_decl_column(
-                api, stmt, node, left_hand_explicit_type, stmt.rvalue
-            )
-        elif type_id is names.RELATIONSHIP:
-            python_type_for_type = infer._infer_type_from_relationship(
-                api, stmt, node, left_hand_explicit_type
-            )
-        elif type_id is names.COLUMN_PROPERTY:
-            python_type_for_type = infer._infer_type_from_decl_column_property(
-                api, stmt, node, left_hand_explicit_type
-            )
-        elif type_id is names.SYNONYM_PROPERTY:
-            python_type_for_type = infer._infer_type_from_left_hand_type_only(
-                api, node, left_hand_explicit_type
-            )
-        elif type_id is names.COMPOSITE_PROPERTY:
-            python_type_for_type = (
-                infer._infer_type_from_decl_composite_property(
-                    api, stmt, node, left_hand_explicit_type
-                )
-            )
-        else:
+        if python_type_for_type is None:
             return
 
     else:
diff --git a/lib/sqlalchemy/ext/mypy/infer.py b/lib/sqlalchemy/ext/mypy/infer.py
index d734d588e..2fea6d340 100644
--- a/lib/sqlalchemy/ext/mypy/infer.py
+++ b/lib/sqlalchemy/ext/mypy/infer.py
@@ -35,6 +35,44 @@ from . import names
 from . import util
 
 
+def _infer_type_from_right_hand_nameexpr(
+    api: SemanticAnalyzerPluginInterface,
+    stmt: AssignmentStmt,
+    node: Var,
+    left_hand_explicit_type: Optional[ProperType],
+    infer_from_right_side: NameExpr,
+) -> Optional[ProperType]:
+
+    type_id = names._type_id_for_callee(infer_from_right_side)
+
+    if type_id is None:
+        return None
+    elif type_id is names.COLUMN:
+        python_type_for_type = _infer_type_from_decl_column(
+            api, stmt, node, left_hand_explicit_type
+        )
+    elif type_id is names.RELATIONSHIP:
+        python_type_for_type = _infer_type_from_relationship(
+            api, stmt, node, left_hand_explicit_type
+        )
+    elif type_id is names.COLUMN_PROPERTY:
+        python_type_for_type = _infer_type_from_decl_column_property(
+            api, stmt, node, left_hand_explicit_type
+        )
+    elif type_id is names.SYNONYM_PROPERTY:
+        python_type_for_type = _infer_type_from_left_hand_type_only(
+            api, node, left_hand_explicit_type
+        )
+    elif type_id is names.COMPOSITE_PROPERTY:
+        python_type_for_type = _infer_type_from_decl_composite_property(
+            api, stmt, node, left_hand_explicit_type
+        )
+    else:
+        return None
+
+    return python_type_for_type
+
+
 def _infer_type_from_relationship(
     api: SemanticAnalyzerPluginInterface,
     stmt: AssignmentStmt,
@@ -255,7 +293,11 @@ def _infer_type_from_decl_column_property(
         # argument
         if type_id is names.COLUMN:
             return _infer_type_from_decl_column(
-                api, stmt, node, left_hand_explicit_type, first_prop_arg
+                api,
+                stmt,
+                node,
+                left_hand_explicit_type,
+                right_hand_expression=first_prop_arg,
             )
 
     return _infer_type_from_left_hand_type_only(
@@ -268,7 +310,7 @@ def _infer_type_from_decl_column(
     stmt: AssignmentStmt,
     node: Var,
     left_hand_explicit_type: Optional[ProperType],
-    right_hand_expression: CallExpr,
+    right_hand_expression: Optional[CallExpr] = None,
 ) -> Optional[ProperType]:
     """Infer the type of mapping from a Column.
 
@@ -305,6 +347,12 @@ def _infer_type_from_decl_column(
 
     callee = None
 
+    if right_hand_expression is None:
+        if not isinstance(stmt.rvalue, CallExpr):
+            return None
+
+        right_hand_expression = stmt.rvalue
+
     for column_arg in right_hand_expression.args[0:2]:
         if isinstance(column_arg, CallExpr):
             if isinstance(column_arg.callee, RefExpr):