Skip to content

Contextually type comprehension arguments during overload resolution#4224

Open
markselby9 wants to merge 1 commit into
facebook:mainfrom
markselby9:feat/4167-overload-comprehension
Open

Contextually type comprehension arguments during overload resolution#4224
markselby9 wants to merge 1 commit into
facebook:mainfrom
markselby9:feat/4167-overload-comprehension

Conversation

@markselby9

Copy link
Copy Markdown
Contributor

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.

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`, so `[(name, "ascending") for name in names]` was
inferred as `list[tuple[str, str]]` and failed to match a
`Sequence[tuple[str, Literal["ascending", "descending"]]]` overload, even
though mypy, pyright, and ty all accept it.

Include comprehension and generator expressions in the un-flattened arm of
`CallWithTypes::type_or_expr`. They 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.

Closes facebook#4167
@github-actions

Copy link
Copy Markdown

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]

scikit-learn (https://github.com/scikit-learn/scikit-learn)
- 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]

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]

optuna (https://github.com/optuna/optuna)
- 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]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Overload resolution doesn't do "expected type" literal narrowing through a comprehension

1 participant