Contextually type comprehension arguments during overload resolution

[markselby9]

Description

Fixes overload resolution failing to contextually type comprehension and generator arguments, so a literal element is not narrowed to the expected type.

When resolving an overloaded call, container literals (list/set/dict) are kept as expressions rather than eagerly inferred, so each candidate overload can contextually type them against its parameter (and narrow element literals). Comprehensions and generator expressions were not — they fell through to a context-free expr_infer in CallWithTypes::type_or_expr, so:

from collections.abc import Sequence
from typing import Literal, overload

Order = Literal["ascending", "descending"]

@overload
def sort_indices(sort_keys: Sequence[tuple[str, Order]]) -> str: ...
@overload
def sort_indices(sort_keys: str) -> int: ...
def sort_indices(sort_keys: Sequence[tuple[str, Order]] | str) -> str | int:
    return 0

names = ["a", "b"]
result = sort_indices([(name, "ascending") for name in names])

was inferred as list[tuple[str, str]] and reported no-matching-overload, even though mypy, pyright, and ty all accept it (reported against Narwhals).

The fix adds Expr::ListComp | Expr::SetComp | Expr::DictComp | Expr::Generator to the same un-flattened arm. These variants already support hint-directed inference via infer_with_decomposed_hint / decompose_list, so deferring their inference lets the existing per-overload contextual-typing path narrow their element/key/value types exactly as it does for container literals. When a hint doesn't decompose (e.g. a generator against a non-Sequence param), inference falls back to the previous context-free behaviour, so no path is worse than before.

Closes #4167

Test plan

Added test_overload_literal_narrowing_through_comprehension (the reported repro) and test_overload_literal_narrowing_through_set_dict_generator (covering the set/dict-comprehension and generator variants) to pyrefly/lib/test/overload.rs. Both fail before the change (no-matching-overload, reveal_type = Unknown) and pass after (reveal_type = str).

Full cargo test -p pyrefly --lib suite passes (0 failures); test.py formatting + linting clean.

[github-actions]

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

rotki (https://github.com/rotki/rotki)
- ERROR rotkehlchen/api/services/assets.py:532:32-535:15: No matching overload found for function `typing.MutableMapping.update` called with arguments: (dict[Asset, list[Price | int]]) [no-matching-overload]

jax (https://github.com/google/jax)
- ERROR jax/experimental/mosaic/gpu/utils.py:845:23-69: The type of `known_factor_prod` is unknown; it is inferred as an implicit `Any` [unknown-variable-type]
- ERROR jax/experimental/roofline/rooflines.py:455:15-63: The type of `axes_size` is unknown; it is inferred as an implicit `Any` [unknown-variable-type]

scikit-learn (https://github.com/scikit-learn/scikit-learn)
- ERROR sklearn/compose/tests/test_column_transformer.py:876:9-35: Result of call expression is of type `ndarray | ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/compose/tests/test_column_transformer.py:876:9-35: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | ndarray` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/compose/tests/test_column_transformer.py:883:9-35: Result of call expression is of type `ndarray | ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/compose/tests/test_column_transformer.py:883:9-35: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | ndarray` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/compose/tests/test_column_transformer.py:2219:9-35: Result of call expression is of type `ndarray | ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/compose/tests/test_column_transformer.py:2219:9-35: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | ndarray` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/compose/tests/test_column_transformer.py:2400:9-35: Result of call expression is of type `ndarray | ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/compose/tests/test_column_transformer.py:2400:9-35: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | ndarray` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/ensemble/_voting.py:430:38-434:14: No matching overload found for function `numpy.lib._shape_base_impl.apply_along_axis` called with arguments: ((x: Unknown) -> signedinteger[_NBitIntP], axis=Literal[1], arr=ndarray[tuple[Any, ...], dtype[Unknown]]) [no-matching-overload]
+ ERROR sklearn/ensemble/_voting.py:430:38-434:14: No matching overload found for function `numpy.lib._shape_base_impl.apply_along_axis` called with arguments: ((x: Unknown) -> signedinteger[_NBitIntP], axis=Literal[1], arr=ndarray[tuple[Any, ...], dtype[Any | Unknown]]) [no-matching-overload]
- ERROR sklearn/ensemble/_voting.py:498:29-37: No matching overload found for function `numpy._core.shape_base.hstack` called with arguments: (ndarray[tuple[Any, ...], dtype[Unknown]]) [no-matching-overload]
+ ERROR sklearn/ensemble/_voting.py:498:29-37: No matching overload found for function `numpy._core.shape_base.hstack` called with arguments: (ndarray[tuple[Any, ...], dtype[Any | Unknown]]) [no-matching-overload]
- ERROR sklearn/ensemble/tests/test_voting.py:110:9-26: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/ensemble/tests/test_voting.py:110:9-26: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/ensemble/tests/test_voting.py:112:9-26: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/ensemble/tests/test_voting.py:112:9-26: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/ensemble/tests/test_voting.py:114:9-28: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/ensemble/tests/test_voting.py:114:9-28: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/metrics/tests/test_classification.py:1038:16-1044:6: The type of `cov_ypyp` is unknown; it is inferred as an implicit `Any` [unknown-variable-type]
- ERROR sklearn/model_selection/_validation.py:414:26-28: Cannot set item in `dict[str, list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]]]` [unsupported-operation]
+ ERROR sklearn/model_selection/_validation.py:414:26-28: Cannot set item in `dict[str, list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]]]` [unsupported-operation]
- ERROR sklearn/tests/test_pipeline.py:1111:9-35: Result of call expression is of type `ndarray | ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/tests/test_pipeline.py:1111:9-35: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | ndarray` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/tests/test_pipeline.py:1369:9-38: Result of call expression is of type `ndarray | ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/tests/test_pipeline.py:1369:9-38: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | ndarray` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/tests/test_pipeline.py:1391:9-38: Result of call expression is of type `ndarray | ndarray[tuple[Any, ...], dtype[Unknown]]` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/tests/test_pipeline.py:1391:9-38: Result of call expression is of type `ndarray[tuple[Any, ...], dtype[Any | Unknown]] | ndarray` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/_metadata_requests.py:1668:29-1677:10: No matching overload found for function `collections.defaultdict.__init__` called with arguments: (type[str], dict[str, None]) [no-matching-overload]
+ ERROR sklearn/utils/_metadata_requests.py:1712:81-87: Identity comparison `None is not '$UNUSED$'` is always True [unnecessary-comparison]
- ERROR sklearn/utils/_param_validation.py:883:47-51: Type of lambda parameter `self` is unknown [implicit-any-lambda]
- ERROR sklearn/utils/tests/test_indexing.py:227:13-61: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/utils/tests/test_indexing.py:227:13-61: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/tests/test_indexing.py:344:13-51: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/utils/tests/test_indexing.py:344:13-51: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/tests/test_indexing.py:367:9-41: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/utils/tests/test_indexing.py:367:9-41: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/tests/test_indexing.py:373:9-49: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/utils/tests/test_indexing.py:373:9-49: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/tests/test_indexing.py:397:9-54: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/utils/tests/test_indexing.py:397:9-54: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/tests/test_indexing.py:405:9-47: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/utils/tests/test_indexing.py:405:9-47: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/tests/test_indexing.py:439:9-43: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
+ ERROR sklearn/utils/tests/test_indexing.py:439:9-43: Result of call expression is of type `list[Unknown] | ndarray[tuple[Any, ...], dtype[Any | Unknown]] | Unknown` and is not used; assign to `_` if this is intentional [unused-call-result]
- ERROR sklearn/utils/validation.py:2575:19-41: No matching overload found for function `numpy._core.fromnumeric.amax` called with arguments: (ndarray[tuple[int], dtype[numpy.bool]] | ndarray[Unknown, Unknown] | None) [no-matching-overload]
+ ERROR sklearn/utils/validation.py:2575:19-41: No matching overload found for function `numpy._core.fromnumeric.amax` called with arguments: (ndarray[tuple[int], dtype[float64]] | ndarray[Unknown, Unknown] | None) [no-matching-overload]
- ERROR sklearn/utils/validation.py:2576:22-44: No matching overload found for function `numpy._core.fromnumeric.amin` called with arguments: (ndarray[tuple[int], dtype[numpy.bool]] | ndarray[Unknown, Unknown] | None) [no-matching-overload]
+ ERROR sklearn/utils/validation.py:2576:22-44: No matching overload found for function `numpy._core.fromnumeric.amin` called with arguments: (ndarray[tuple[int], dtype[float64]] | ndarray[Unknown, Unknown] | None) [no-matching-overload]

apprise (https://github.com/caronc/apprise)
- ERROR apprise/plugins/jira.py:781:22-783:10: No matching overload found for function `typing.MutableMapping.update` called with arguments: (dict[LiteralString, str]) [no-matching-overload]

zulip (https://github.com/zulip/zulip)
+ ERROR zerver/actions/message_send.py:1411:12-67: Returned type `list[object]` is not assignable to declared return type `list[str]` [bad-return]

xarray (https://github.com/pydata/xarray)
- ERROR xarray/core/dataset.py:2747:23-27: Yielded type `tuple[Hashable, _arrayapi[Any, Any] | _arrayfunction[Any, Any] | ndarray | ndarray[Unknown, dtype[Unknown]] | Unknown]` is not assignable to declared yield type `tuple[Hashable, Variable | int | ndarray | slice[Any, Any, Any]]` [invalid-yield]
+ ERROR xarray/core/dataset.py:2747:23-27: Yielded type `tuple[Hashable, _arrayapi[Any, Any] | _arrayfunction[Any, Any] | ndarray | ndarray[Unknown, dtype[Any | Unknown]] | ndarray[Unknown, dtype[Unknown]] | Unknown]` is not assignable to declared yield type `tuple[Hashable, Variable | int | ndarray | slice[Any, Any, Any]]` [invalid-yield]

optuna (https://github.com/optuna/optuna)
+ ERROR optuna/samplers/_tpe/sampler.py:901:19-34: The type of `pareto_sols` is unknown; it is inferred as an implicit `Any` [unknown-variable-type]
- ERROR optuna/visualization/_hypervolume_history.py:114:30-50: `*` is not supported between `ndarray[tuple[int], dtype[signedinteger[_NBitIntP]]]` and `None` [unsupported-operation]
+ ERROR optuna/visualization/_hypervolume_history.py:114:30-50: `*` is not supported between `ndarray[tuple[int], dtype[float64]]` and `None` [unsupported-operation]