Use type context for for loops

[A5rocks]

Fixes #18826. This PR makes for loops use any annotations on the index variable for type context when checking the thing being iterated over.

I'm not confident this should be merged. This breaks this:

from typing import Literal

def func() -> None:
    b: None | str
    for b in ["a", "c"]:
        b + "5"

I think it should narrow somehow but I'm not quite sure how to do that.

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

speedrun.com_global_scoreboard_webapp (https://github.com/Avasam/speedrun.com_global_scoreboard_webapp)
+ backend/api/tournament_scheduler_api.py:291: error: Need type annotation for "time_slot" (hint: "time_slot: list[<type>] = ...")  [var-annotated]

cwltool (https://github.com/common-workflow-language/cwltool)
+ tests/test_js_sandbox.py: note: In function "hide_nodejs":
+ tests/test_js_sandbox.py:69:25: error: Argument 1 to "Path" has incompatible type "str | None"; expected "str | PathLike[str]"  [arg-type]
+ tests/test_js_sandbox.py:70:34: error: Argument 1 to "remove" of "list" has incompatible type "str | None"; expected "str"  [arg-type]

pandas (https://github.com/pandas-dev/pandas)
+ pandas/io/common.py:1161: error: Unused "type: ignore" comment  [unused-ignore]
+ pandas/io/common.py:1161: error: Item "str" of "str | BaseBuffer" has no attribute "close"  [union-attr]
+ pandas/io/common.py:1161: note: Error code "union-attr" not covered by "type: ignore" comment
+ pandas/io/common.py:1161: error: Item "BaseBuffer" of "str | BaseBuffer" has no attribute "close"  [union-attr]
- pandas/core/internals/managers.py:972: error: Incompatible types in assignment (expression has type "signedinteger[_32Bit | _64Bit]", variable has type "int")  [assignment]

steam.py (https://github.com/Gobot1234/steam.py)
+ steam/state.py:351: error: Item "None" of "list[int] | None" has no attribute "__iter__" (not iterable)  [union-attr]
- steam/state.py:1747: error: Incompatible types in assignment (expression has type "User", variable has type "int")  [assignment]
+ steam/state.py:1747: error: Argument 2 to "zip" has incompatible type "Sequence[User]"; expected "Iterable[int]"  [arg-type]

spark (https://github.com/apache/spark)
+ python/pyspark/sql/profiler.py:267: error: Argument 1 to "list" has incompatible type "dict_items[int, tuple[Stats | None, dict[str, list[tuple[int, tuple[float, float, int] | None]]] | None]]"; expected "Iterable[tuple[int | None, tuple[Stats | None, dict[str, list[tuple[int, tuple[float, float, int] | None]]] | None, Any]]]"  [arg-type]
+ python/pyspark/sql/profiler.py:268: error: Invalid index type "int | None" for "dict[int, tuple[Stats | None, dict[str, list[tuple[int, tuple[float, float, int] | None]]] | None]]"; expected type "int"  [index]
+ python/pyspark/sql/profiler.py:270: error: Argument 1 to "pop" of "dict" has incompatible type "int | None"; expected "int"  [arg-type]
+ python/pyspark/sql/profiler.py:292: error: Argument 1 to "list" has incompatible type "dict_items[int, tuple[Stats | None, dict[str, list[tuple[int, tuple[float, float, int] | None]]] | None]]"; expected "Iterable[tuple[int | None, tuple[Stats | None, dict[str, list[tuple[int, tuple[float, float, int] | None]]] | None, Any]]]"  [arg-type]
+ python/pyspark/sql/profiler.py:293: error: Invalid index type "int | None" for "dict[int, tuple[Stats | None, dict[str, list[tuple[int, tuple[float, float, int] | None]]] | None]]"; expected type "int"  [index]
+ python/pyspark/sql/profiler.py:295: error: Argument 1 to "pop" of "dict" has incompatible type "int | None"; expected "int"  [arg-type]

sockeye (https://github.com/awslabs/sockeye)
- sockeye/translate.py:190: error: Incompatible types in assignment (expression has type "tuple[Any, ...]", variable has type "list[str]")  [assignment]
+ sockeye/translate.py:190: error: Argument 1 to "enumerate" has incompatible type "zip[tuple[Any, ...]]"; expected "Iterable[list[str]]"  [arg-type]

mypy (https://github.com/python/mypy)
+ mypy/checkpattern.py:312: error: Argument 1 to "conditional_types_with_intersection" of "TypeChecker" has incompatible type "Optional[Type]"; expected "Type"  [arg-type]
+ mypy/checkpattern.py:312: error: Argument "default" to "conditional_types_with_intersection" of "TypeChecker" has incompatible type "Optional[Type]"; expected "None"  [arg-type]
+ mypy/checkpattern.py:317: error: Argument 1 to "is_uninhabited" has incompatible type "Optional[Type]"; expected "Type"  [arg-type]
+ mypy/checkpattern.py:318: error: Argument 1 to "TupleType" has incompatible type "list[Optional[Type]]"; expected "list[Type]"  [arg-type]
+ mypy/checkpattern.py:322: error: Argument 1 to "is_uninhabited" has incompatible type "Optional[Type]"; expected "Type"  [arg-type]
+ mypy/checkpattern.py:325: error: Argument 1 to "is_uninhabited" has incompatible type "Optional[Type]"; expected "Type"  [arg-type]
+ mypy/checkpattern.py:330: error: List comprehension has incompatible type List[Optional[Type]]; expected List[Type]  [misc]
+ mypy/checkpattern.py:330: error: Argument 1 to "is_uninhabited" has incompatible type "Optional[Type]"; expected "Type"  [arg-type]
+ mypyc/irbuild/ll_builder.py:2324: error: "RType" has no attribute "class_ir"  [attr-defined]

xarray (https://github.com/pydata/xarray)
- xarray/core/indexing.py: note: In function "_decompose_outer_indexer":
- xarray/core/indexing.py:1260: error: Argument 1 to "_decompose_slice" has incompatible type "number[Any, int | float | complex]"; expected "slice[Any, Any, Any]"  [arg-type]
- xarray/core/indexing.py:1298: error: Argument 1 to "_decompose_slice" has incompatible type "number[Any, int | float | complex]"; expected "slice[Any, Any, Any]"  [arg-type]
- xarray/core/indexing.py: note: At top level:
+ xarray/tests/test_backends.py:6227: error: Unused "type: ignore" comment  [unused-ignore]

porcupine (https://github.com/Akuli/porcupine)
+ porcupine/plugins/directory_tree.py:298: error: Item "None" of "Path | None" has no attribute "is_dir"  [union-attr]
+ porcupine/plugins/directory_tree.py:303: error: Item "None" of "Path | None" has no attribute "name"  [union-attr]
+ porcupine/plugins/directory_tree.py:305: error: Item "None" of "Path | None" has no attribute "is_dir"  [union-attr]

freqtrade (https://github.com/freqtrade/freqtrade)
+ freqtrade/optimize/backtesting.py:1436: error: Item "None" of "LocalTrade | None" has no attribute "select_order"  [union-attr]
+ freqtrade/optimize/backtesting.py:1436: error: Item "None" of "LocalTrade | None" has no attribute "entry_side"  [union-attr]
+ freqtrade/optimize/backtesting.py:1437: error: Argument 2 to "_try_close_open_order" of "Backtesting" has incompatible type "LocalTrade | None"; expected "LocalTrade"  [arg-type]
+ freqtrade/optimize/backtesting.py:1441: error: Item "None" of "LocalTrade | None" has no attribute "has_open_position"  [union-attr]
+ freqtrade/optimize/backtesting.py:1442: error: Argument 1 to "_check_trade_exit" of "Backtesting" has incompatible type "LocalTrade | None"; expected "LocalTrade"  [arg-type]
+ freqtrade/optimize/backtesting.py:1445: error: Item "None" of "LocalTrade | None" has no attribute "select_order"  [union-attr]
+ freqtrade/optimize/backtesting.py:1445: error: Item "None" of "LocalTrade | None" has no attribute "exit_side"  [union-attr]
+ freqtrade/optimize/backtesting.py:1447: error: Argument 2 to "_process_exit_order" of "Backtesting" has incompatible type "LocalTrade | None"; expected "LocalTrade"  [arg-type]

jinja (https://github.com/pallets/jinja)
+ src/jinja2/debug.py:70: error: Item "None" of "Optional[TracebackType]" has no attribute "tb_next"  [union-attr]

materialize (https://github.com/MaterializeInc/materialize)
- misc/python/materialize/cli/cloudbench.py:173: error: Incompatible types in assignment (expression has type "Generator[dict_values[str | Any, str | Any], None, None]", variable has type "BenchSuccessResult | BenchFailureLogs | None")  [assignment]
+ misc/python/materialize/cli/cloudbench.py:173: error: Argument 2 to "zip" has incompatible type "Generator[Generator[dict_values[str | Any, str | Any], None, None], None, None]"; expected "Iterable[BenchSuccessResult | BenchFailureLogs | None]"  [arg-type]

pylint (https://github.com/pycqa/pylint)
+ pylint/checkers/strings.py:399: error: Argument 2 to "arg_matches_format_type" has incompatible type "str | None"; expected "str"  [arg-type]

Tanjun (https://github.com/FasterSpeeding/Tanjun)
+ tanjun/_internal/__init__.py:724: error: Item "list[tuple[list[str], MessageCommand[Any]]]" of "_TreeT | list[tuple[list[str], MessageCommand[Any]]]" has no attribute "get"  [union-attr]
- tanjun/clients.py:2294: error: "Never" has no attribute "__iter__" (not iterable)  [attr-defined]
- tanjun/clients.py:2294: error: Argument 1 to "call_with_async_di" of "Context" has incompatible type "Callable[[MessageContext, VarArg(Any), KwArg(Any)], Coroutine[Any, Any, Iterable[str]]]"; expected "Callable[..., Coroutine[Any, Any, Never] | Never]"  [arg-type]

jax (https://github.com/google/jax)
+ jax/experimental/mosaic/gpu/dialect_lowering.py:696: error: INTERNAL ERROR -- Please try using mypy master on GitHub:
+ https://mypy.readthedocs.io/en/stable/common_issues.html#using-a-development-mypy-build
+ Please report a bug at https://github.com/python/mypy/issues
+ version: 1.16.0+dev.bef1077d9fe4887e86f958f2bcbd1ab83424fbe1
+ jax/experimental/mosaic/gpu/dialect_lowering.py:696: : note: use --pdb to drop into pdb
- jax/experimental/jax2tf/jax2tf.py:452: error: Unused "type: ignore" comment  [unused-ignore]
- jax/experimental/jax2tf/jax2tf.py:3552: error: Unused "type: ignore" comment  [unused-ignore]
- jax/experimental/jax2tf/call_tf.py:580: error: Unused "type: ignore" comment  [unused-ignore]
- jax/_src/pallas/core.py:1176: error: Unused "type: ignore" comment  [unused-ignore]
- jax/_src/pallas/mosaic_gpu/core.py:238: error: Variable "jax._src.pallas.mosaic_gpu.core.Index" is not valid as a type  [valid-type]
- jax/_src/pallas/mosaic_gpu/core.py:238: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
- jax/_src/pallas/mosaic_gpu/core.py:239: error: Variable "jax._src.pallas.mosaic_gpu.core.Index" is not valid as a type  [valid-type]
- jax/_src/pallas/mosaic_gpu/core.py:239: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
- jax/_src/pallas/mosaic_gpu/core.py:242: error: Variable "jax._src.pallas.mosaic_gpu.core.Index" is not valid as a type  [valid-type]
- jax/_src/pallas/mosaic_gpu/core.py:242: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
- jax/_src/pallas/mosaic_gpu/core.py:327: error: Variable "jax._src.pallas.mosaic_gpu.core.Index" is not valid as a type  [valid-type]
- jax/_src/pallas/mosaic_gpu/core.py:327: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
- jax/_src/pallas/mosaic_gpu/core.py:328: error: Variable "jax._src.pallas.mosaic_gpu.core.Index" is not valid as a type  [valid-type]
- jax/_src/pallas/mosaic_gpu/core.py:328: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
- jax/_src/pallas/mosaic_gpu/core.py:399: error: Variable "jax._src.pallas.mosaic_gpu.core.Index" is not valid as a type  [valid-type]
- jax/_src/pallas/mosaic_gpu/core.py:399: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
- jax/_src/pallas/mosaic_gpu/core.py:400: error: Variable "jax._src.pallas.mosaic_gpu.core.Index" is not valid as a type  [valid-type]
- jax/_src/pallas/mosaic_gpu/core.py:400: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
- jax/_src/pallas/mosaic/verification.py:553: error: Unused "type: ignore" comment  [unused-ignore]
- jax/_src/pallas/mosaic_gpu/primitives.py:129: error: Unused "type: ignore" comment  [unused-ignore]
+ Traceback (most recent call last):
+   File "", line 8, in <module>
+     sys.exit(console_entry())
+   File "/__main__.py", line 15, in console_entry
+     main()
+   File "/main.py", line 126, in main
+     res, messages, blockers = run_build(sources, options, fscache, t0, stdout, stderr)
+   File "/main.py", line 210, in run_build
+     res = build.build(sources, options, None, flush_errors, fscache, stdout, stderr)
+   File "/build.py", line 191, in build
+     result = _build(
+   File "/build.py", line 267, in _build
+     graph = dispatch(sources, manager, stdout)
+   File "/build.py", line 2939, in dispatch
+     process_graph(graph, manager)
+   File "/build.py", line 3337, in process_graph
+     process_stale_scc(graph, scc, manager)
+   File "/build.py", line 3438, in process_stale_scc
+     graph[id].type_check_first_pass()
+   File "/build.py", line 2311, in type_check_first_pass
+     self.type_checker().check_first_pass()
+   File "/checker.py", line 474, in check_first_pass
+     self.accept(d)
+   File "/checker.py", line 581, in accept
+     stmt.accept(self)
+   File "/nodes.py", line 1462, in accept
+     return visitor.visit_for_stmt(self)
+   File "/checker.py", line 5321, in visit_for_stmt
+     iterator_type, item_type = self.analyze_iterable_item_type(s.expr, context)
+   File "/checker.py", line 5351, in analyze_iterable_item_type
+     self.expr_checker.accept(expr, context), context=expr
+   File "/checkexpr.py", line 5954, in accept
+     typ = node.accept(self)
+   File "/nodes.py", line 2319, in accept
+     return visitor.visit_list_expr(self)
+            ~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
+   File "/checkexpr.py", line 5034, in visit_list_expr
+     return self.check_lst_expr(e, "builtins.list", "<list>")
+            ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   File "/checkexpr.py", line 5100, in check_lst_expr
+     out = self.check_call(
+           ~~~~~~~~~~~~~~~^
+         constructor,
+         ^^^^^^^^^^^^
+     ...<2 lines>...
+         e,
+         ^^
+     )[0]
+     ^
+   File "/checkexpr.py", line 1579, in check_call
+     return self.check_callable_call(
+            ~~~~~~~~~~~~~~~~~~~~~~~~^
+         callee,
+         ^^^^^^^
+     ...<6 lines>...
+         object_type,
+         ^^^^^^^^^^^^
+     )
+     ^
+   File "/checkexpr.py", line 1763, in check_callable_call
+     callee = self.infer_function_type_arguments(
+         callee, args, arg_kinds, arg_names, formal_to_actual, need_refresh, context
+     )
+   File "/checkexpr.py", line 2094, in infer_function_type_arguments
+     inferred_args, _ = infer_function_type_arguments(
+                        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
+         callee_type,
+         ^^^^^^^^^^^^
+     ...<5 lines>...
+         strict=self.chk.in_checked_function(),
+         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+     )
+     ^
+   File "/infer.py", line 57, in infer_function_type_arguments
+     constraints = infer_constraints_for_callable(
+         callee_type, arg_types, arg_kinds, arg_names, formal_to_actual, context
+     )
+   File "/constraints.py", line 252, in infer_constraints_for_callable
+     c = infer_constraints(callee.arg_types[i], actual_type, SUPERTYPE_OF)
+   File "/constraints.py", line 320, in infer_constraints
+     return _infer_constraints(template, actual, direction, skip_neg_op)
+   File "/constraints.py", line 428, in _infer_constraints
+     return template.accept(ConstraintBuilderVisitor(actual, direction, skip_neg_op))
+            ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   File "/types.py", line 2458, in accept
+     return visitor.visit_tuple_type(self)
+            ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
+   File "/constraints.py", line 1318, in visit_tuple_type
+     infer_constraints(template_items[i], actual_items[i], self.direction)
+     ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   File "/constraints.py", line 320, in infer_constraints
+     return _infer_constraints(template, actual, direction, skip_neg_op)
+   File "/constraints.py", line 428, in _infer_constraints
+     return template.accept(ConstraintBuilderVisitor(actual, direction, skip_neg_op))
+            ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   File "/types.py", line 3039, in accept
+     return visitor.visit_partial_type(self)
+            ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
+   File "/constraints.py", line 706, in visit_partial_type
+     assert False, "Internal error"
+            ^^^^^
+ AssertionError: Internal error

[A5rocks]

Yeah that's not a good mypy primer output

[A5rocks]

I think this needs to copy the new redefinition rules and therefore I'll close this for now.