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, 50 insertions, 2 deletions

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):