Fix crashes and fails in forward references

mypy/semanal.py

@@ -4230,6 +4230,14 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
                     self.analyze(node.type_override, node)
         super().visit_assignment_stmt(s)
 
+    def visit_for_stmt(self, s: ForStmt) -> None:
+        self.analyze(s.index_type, s)
+        super().visit_for_stmt(s)
+
+    def visit_with_stmt(self, s: WithStmt) -> None:
+        self.analyze(s.target_type, s)
+        super().visit_with_stmt(s)
+
     def visit_cast_expr(self, e: CastExpr) -> None:
         self.analyze(e.type, e)
         super().visit_cast_expr(e)
@@ -4247,6 +4255,14 @@ def visit_type_application(self, e: TypeApplication) -> None:
     def perform_transform(self, node: Union[Node, SymbolTableNode],
                           transform: Callable[[Type], Type]) -> None:
         """Apply transform to all types associated with node."""
+        if isinstance(node, ForStmt):
+            node.index_type = transform(node.index_type)
+            self.transform_types(node.index, transform)
+        if isinstance(node, WithStmt):
+            node.target_type = transform(node.target_type)
+            for n in node.target:
+                if isinstance(n, NameExpr) and isinstance(n.node, Var) and n.node.type:
+                    n.node.type = transform(n.node.type)
         if isinstance(node, (FuncDef, CastExpr, AssignmentStmt, TypeAliasExpr, Var)):
             node.type = transform(node.type)
         if isinstance(node, NewTypeExpr):
@@ -4275,6 +4291,15 @@ def perform_transform(self, node: Union[Node, SymbolTableNode],
                     new_bases.append(alt_base)
             node.bases = new_bases
 
+    def transform_types(self, lvalue: Lvalue, transform: Callable[[Type], Type]) -> None:
+        if isinstance(lvalue, RefExpr):
+            if isinstance(lvalue.node, Var):
+                var = lvalue.node
+                var.type = transform(var.type)
+        elif isinstance(lvalue, TupleExpr):
+            for item in lvalue.items:
+                self.transform_types(item, transform)
+
     def analyze(self, type: Optional[Type], node: Union[Node, SymbolTableNode],
                 warn: bool = False) -> None:
         # Recursive type warnings are only emitted on type definition 'node's, marked by 'warn'

test-data/unit/check-statements.test

@@ -1565,25 +1565,51 @@ reveal_type(d['weight0']) # E: Revealed type is 'builtins.float*'
 
 [builtins fixtures/floatdict.pyi]
 
-[case testForwardRefsInForStatement]
+[case testForwardRefsInForStatementImplicit]
 from typing import List, NamedTuple
 lst: List[N]
 
+for i in lst:
+    a: int = i.x
+    b: str = i[0] # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+
+N = NamedTuple('N', [('x', int)])
+[builtins fixtures/list.pyi]
+[out]
+
+[case testForwardRefsInForStatement]
+from typing import List, NamedTuple
+lst: List[M]
+
 for i in lst: # type: N
     a: int = i.x
-    b: str = i[0]
+    b: str = i[0] # E: Incompatible types in assignment (expression has type "int", variable has type "str")
 
 N = NamedTuple('N', [('x', int)])
+class M(N): pass
 [builtins fixtures/list.pyi]
 [out]
 
-[case testForwardRefsInWithStatement]
-from typing import ContextManager
+[case testForwardRefsInWithStatementImplicit]
+from typing import ContextManager, Any
 from mypy_extensions import TypedDict
 cm: ContextManager[N]
 
+with cm as g:
+    a: int = g['x']
+
+N = TypedDict('N', {'x': int})
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+[out]
+
+[case testForwardRefsInWithStatement]
+from typing import ContextManager, Any
+from mypy_extensions import TypedDict
+cm: ContextManager[Any]
+
 with cm as g:  # type: N
-    a: str = g['x']
+    a: str = g['x']  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
 
 N = TypedDict('N', {'x': int})
 [builtins fixtures/dict.pyi]