Fix bidirectional type checking not working for lambdas

[hrolfurgylfa]

Summary

I addressed the missing bidirectional type-checking for lambdas with a fairly easy one line fix, but it revealed a separate problem with overloads leaking errors between each other in my sort unit test so I ended up fixing that as the first commit with a separate regression test.

There is one other unit test that got modified by this change, as it was relying on the previous behvaiour of eagerly inferring lambdas in function calls, but it should be fine this way as well. Pyrefly should never have a problem with casting a literal back down to a string right?

Fixes #3950

Test Plan

Only tested through unit tests, this only touches on the solving step, so it shouldn't be close to anything else.

No files generated by test.py

[hrolfurgylfa]

This is a lot of changes through mypy_primer, I might split the two commits up into separate PRs to make it easier to go through each one individually.

[hrolfurgylfa]

Nevermind, this is all the same bug, this overload error propagation doesn't happen anywhere else. I can't get it to happen with mutable builtin containers either, something like this works just fine even though I'd expect it to exercise the exact same path. I guess I might be fixing this in the wrong place and these containers address the same problem elsewhere.

But for now I'll just go through mypy_primer results one by one and see if they are all expected. Will re-open the PR once I have reviewed all of those.

[github-actions]

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

egglog-python (https://github.com/egraphs-good/egglog-python)
- ERROR python/egglog/declarations.py:172:65-174:6: `dict[Ident, RulesetDecl]` is not assignable to `dict[Ident, CombinedRulesetDecl | RulesetDecl]` [bad-assignment]

dulwich (https://github.com/dulwich/dulwich)
- ERROR dulwich/diff_tree.py:395:61-397:6: `defaultdict[Unknown, list[None]]` is not assignable to `dict[bytes, list[TreeChange | None]]` [bad-assignment]

rotki (https://github.com/rotki/rotki)
- ERROR rotkehlchen/mcp/registry.py:48:18-51:6: No matching overload found for function `sorted` called with arguments: (list[(...) -> Any], key=(tool: Unknown) -> Any | None) [no-matching-overload]

steam.py (https://github.com/Gobot1234/steam.py)
- ERROR steam/chat.py:261:28-267:14: No matching overload found for function `steam.reaction.MessageReaction.__init__` called with arguments: (ConnectionState, Self@ChatMessage, Emoticon | None, Sticker | None, user=ClanMember | GroupMember | PartialMember) [no-matching-overload]
+ ERROR steam/chat.py:261:28-267:14: No matching overload found for function `steam.reaction.MessageReaction.__init__` called with arguments: (ConnectionState, Self@ChatMessage, Unknown, Unknown, user=ClanMember | GroupMember | PartialMember) [no-matching-overload]

jax (https://github.com/google/jax)
+ ERROR jax/experimental/jax2tf/tests/tf_test_util.py:238:50-63: Object of class `object` has no attribute `skip_tf_run` [missing-attribute]
+ ERROR jax/experimental/jax2tf/tests/tf_test_util.py:249:54-71: Object of class `object` has no attribute `expect_tf_error` [missing-attribute]
+ ERROR jax/experimental/jax2tf/tests/tf_test_util.py:269:54-71: Object of class `object` has no attribute `skip_comparison` [missing-attribute]
+ ERROR jax/experimental/jax2tf/tests/tf_test_util.py:275:52-97: Object of class `object` has no attribute `get_max_tolerance_limitation` [missing-attribute]
+ ERROR jax/experimental/jax2tf/tests/tf_test_util.py:283:56-71: Object of class `object` has no attribute `custom_assert` [missing-attribute]
+ ERROR jax/experimental/jax2tf/tests/tf_test_util.py:293:11-45: Object of class `object` has no attribute `custom_assert` [missing-attribute]

spark (https://github.com/apache/spark)
+ ERROR python/pyspark/sql/conversion.py:1840:38-48: Cannot index into `NAType` [bad-index]
+ ERROR python/pyspark/sql/conversion.py:1840:50-63: Cannot index into `NAType` [bad-index]
+ ERROR python/pyspark/sql/conversion.py:1845:45-53: Cannot index into `NAType` [bad-index]
+ ERROR python/pyspark/sql/conversion.py:1845:55-64: Cannot index into `NAType` [bad-index]
+ ERROR python/pyspark/sql/conversion.py:1850:44-52: Cannot index into `NAType` [bad-index]
+ ERROR python/pyspark/sql/conversion.py:1850:54-63: Cannot index into `NAType` [bad-index]

prefect (https://github.com/PrefectHQ/prefect)
- ERROR src/prefect/cli/artifact.py:96:45-63: Object of class `Artifact` has no attribute `latest_id` [missing-attribute]
+ ERROR src/prefect/cli/artifact.py:93:31-63: No matching overload found for function `sorted` called with arguments: (list[Artifact] | list[ArtifactCollection], key=(x: Unknown) -> Unknown) [no-matching-overload]
+ ERROR src/prefect/cli/flow_run.py:265:31-81: No matching overload found for function `sorted` called with arguments: (list[FlowRun], key=(d: Unknown) -> Unknown, reverse=Literal[True]) [no-matching-overload]
- ERROR src/prefect/cli/flow_run.py:268:17-45: Object of class `NoneType` has no attribute `state_details` [missing-attribute]
- ERROR src/prefect/cli/flow_run.py:269:20-47: Object of class `NoneType` has no attribute `is_scheduled` [missing-attribute]
- ERROR src/prefect/cli/flow_run.py:270:22-46: Object of class `NoneType` has no attribute `timestamp` [missing-attribute]
- ERROR src/prefect/cli/flow_run.py:276:21-40: Object of class `NoneType` has no attribute `type` [missing-attribute]
+ ERROR src/prefect/server/services/task_run_recorder.py:225:39-69: Object of class `NoneType` has no attribute `timestamp` [missing-attribute]

freqtrade (https://github.com/freqtrade/freqtrade)
+ ERROR freqtrade/data/btanalysis/trade_parallelism.py:124:51-78: Cannot index into `int` [bad-index]
+ ERROR freqtrade/data/btanalysis/trade_parallelism.py:124:51-78: Cannot index into `str` [bad-index]
+ ERROR freqtrade/plugins/pairlist/PercentChangePairList.py:233:32-237:10: No matching overload found for function `sorted` called with arguments: (list[SymbolWithPercentage], reverse=bool, key=(t: Unknown) -> Unknown) [no-matching-overload]

pwndbg (https://github.com/pwndbg/pwndbg)
+ ERROR pwndbg/aglib/disasm/aarch64.py:331:75-86: Expected 0 positional arguments, got 2 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/aarch64.py:535:17-23: Expected 0 positional arguments, got 3 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/arm.py:267:75-86: Expected 0 positional arguments, got 2 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/arm.py:435:17-23: Expected 0 positional arguments, got 3 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/assistant.py:345:78-89: Expected 0 positional arguments, got 3 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/assistant.py:380:21-32: Expected 0 positional arguments, got 3 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/assistant.py:392:66-77: Expected 0 positional arguments, got 2 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/mips.py:233:75-86: Expected 0 positional arguments, got 2 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/riscv.py:199:75-86: Expected 0 positional arguments, got 2 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/sparc.py:103:80-83: Expected 0 positional arguments, got 1 [bad-argument-count]
+ ERROR pwndbg/aglib/disasm/x86.py:153:75-86: Expected 0 positional arguments, got 2 [bad-argument-count]
+ ERROR pwndbg/enhance.py:174:16-25: Returned type `object` is not assignable to declared return type `str` [bad-return]
+ ERROR pwndbg/enhance.py:176:62-74: No matching overload found for function `str.join` called with arguments: (tuple[object, ...]) [no-matching-overload]

schemathesis (https://github.com/schemathesis/schemathesis)
- ERROR src/schemathesis/cli/commands/run/handlers/output.py:597:70-599:6: `dict[PhaseName, tuple[Status, None]]` is not assignable to `dict[PhaseName, tuple[Status, PhaseSkipReason | None]]` [bad-assignment]

ibis (https://github.com/ibis-project/ibis)
+ ERROR ibis/backends/conftest.py:268:66-68: Missing argument `backend` [missing-argument]

strawberry (https://github.com/strawberry-graphql/strawberry)
+ ERROR strawberry/federation/schema.py:412:55-82: Generic attribute `__strawberry_definition__` of class `WithStrawberryDefinition` is not visible on the class
+ Object of class `StrawberryType` has no attribute `__strawberry_definition__` [missing-attribute]

zulip (https://github.com/zulip/zulip)
+ ERROR zerver/lib/export.py:3209:18-84: No matching overload found for function `sorted` called with arguments: (dict_values[int, dict[str, bool | int | str | None]], key=(export_dict: Unknown) -> Unknown) [no-matching-overload]

sockeye (https://github.com/awslabs/sockeye)
- ERROR sockeye_contrib/plot_metrics.py:24:19-33:3: No matching overload found for function `typing.MutableMapping.update` called with arguments: (dict[str, ((s: Unknown) -> float) | type[float] | type[int]]) [no-matching-overload]

meson (https://github.com/mesonbuild/meson)
- ERROR mesonbuild/cargo/interpreter.py:99:57-101:6: `PerMachine[None]` is not assignable to `PerMachine[PackageConfiguration | None]` [bad-assignment]
- ERROR mesonbuild/cargo/manifest.py:375:38-88: `list[str]` is not assignable to `list[CRATE_TYPE]` [bad-assignment]
+ ERROR mesonbuild/dependencies/cuda.py:127:75-79: Argument `str | None` is not assignable to parameter `s` with type `str` in function `mesonbuild.utils.universal.Version.__init__` [bad-argument-type]

werkzeug (https://github.com/pallets/werkzeug)
+ ERROR tests/test_exceptions.py:142:11-145:6: No matching overload found for function `sorted` called with arguments: (Generator[type[HTTPException]], key=(e: Unknown) -> Unknown) [no-matching-overload]
+ ERROR tests/test_exceptions.py:162:11-165:6: No matching overload found for function `sorted` called with arguments: (Generator[type[HTTPException]], key=(e: Unknown) -> Unknown) [no-matching-overload]

core (https://github.com/home-assistant/core)
- ERROR homeassistant/components/conversation/chat_log.py:339:30-88: `list[SystemContent]` is not assignable to `list[Content]` [bad-assignment]
- ERROR homeassistant/components/freebox/entity.py:105:16-26: Returning implicit Any from function declared to return "int | None" [no-any-return-implicit]

discord.py (https://github.com/Rapptz/discord.py)
+ ERROR discord/channel.py:3411:47-51: Object of class `NoneType` has no attribute `id` [missing-attribute]
+ ERROR discord/guild.py:1283:41-47: Object of class `NoneType` has no attribute `name` [missing-attribute]
+ ERROR discord/message.py:576:21-33: Object of class `NoneType` has no attribute `created_at` [missing-attribute]
+ ERROR discord/state.py:525:37-41: Object of class `NoneType` has no attribute `id` [missing-attribute]

optuna (https://github.com/optuna/optuna)
+ ERROR tutorial/20_recipes/002_multi_objective.py:150:68-79: `None` is not subscriptable [unsupported-operation]
+ ERROR optuna/samplers/nsgaii/_elite_population_selection_strategy.py:85:39-50: `None` is not subscriptable [unsupported-operation]

schema_salad (https://github.com/common-workflow-language/schema_salad)
+ ERROR src/schema_salad/dlang_codegen.py:440:50-59: `&` is not supported between `Sized` and `Sized` [unsupported-operation]

[hrolfurgylfa]

I've gone through about half of the mypy_primer list, here are my thoughts so far:

First 3 are just fixes because of bidirectional type checking

steam.py (https://github.com/Gobot1234/steam.py)

This is using a function similar to map/filter, the first argument is the lambda and the second argument is the iterable that has the type variable information. The reason this regressed seems to be that previously the lambda would be solved to Unknown and the function would use the second argument to infer the type, but now that the lambda isn't eagerly evaluated, it now tries to use the lambda to find the type, which fails into it being Unknown. Not completely sure on this one though, but it would be solved by #438

Here is a minimal example of this regression:

from collections.abc import AsyncIterable, Awaitable, Callable, Iterable
from dataclasses import dataclass
from typing import overload, reveal_type


@overload
def find[T](predicate: Callable[[T], bool], iterable: AsyncIterable[T], /) -> Awaitable[T | None]: ...


@overload
def find[T](predicate: Callable[[T], bool], iterable: Iterable[T], /) -> T | None: ...


def find[T](
    predicate: Callable[[T], bool], iterable: Iterable[T] | AsyncIterable[T], /
) -> T | None | Awaitable[T | None]: ...


@dataclass
class WithName:
    name: str


@dataclass
class PartialReaction:
    emoticon: WithName | None
    sticker: WithName


emoticon: WithName = WithName("abc")
partial_reactions: list[PartialReaction] = []

reaction = find(lambda r: (r.emoticon or r.sticker).name == emoticon.name, partial_reactions)
reveal_type(reaction)  # E: revealed type: PartialReaction | None

jax (https://github.com/google/jax)

These are all showing up because the lambda return in the following statement is no longer inferred to be Unknown, so it went further in the type checking process and somehow ended up with object. To be fully solved this would require the more complex bidirectional type checking in #438

jax2tf_limits = tuple(filter(lambda l: l.filter(), limitations))

spark (https://github.com/apache/spark)

These errors are now showing up inside lambdas as previously the argument inside them was Unknown, since they were eagerly evaluated.

prefect (https://github.com/PrefectHQ/prefect)

TODO, very confusing why it fails to find the overload of sorted

freqtrade (https://github.com/freqtrade/freqtrade)

trade_parallelism.py:

Error in the lambda of a sorted call now showing, just a fault of actually type checking the lambda now.

PercentChangePairList.py:

Not exactly sure why it isn't being inferred into the second overload of sorted, I thought Unknown would be treated as Any and it would be able to match into it just fine.

pwndbg (https://github.com/pwndbg/pwndbg)

For the first ones:

Lambdas that accept *args seem to be treated as zero argument functions instead of taking any number of arguments, hence getting these errors calling for too few arguments. Reproduction example:

def add(a: int, b: int) -> int:
    return a + b


def mul(a: int, b: int) -> int:
    return a * b


FUNCS = {"add": add, "mul": mul}

key: str = "any"
a = FUNCS.get(key, lambda *c: 3)(1, 2)

For the last two:

A filter(lambda x: x is not None) is run on a list with a really large union list[str | Unknown] | list[str | Unknown | None] | list[str] | list[Unknown] and it collapses into a list[object], not really sure why that is happening.

[hrolfurgylfa]

At this point, I'm not too sure if this is this right direction I am heading in, if it is, I could continue, but seeing if these regressions can be fixed and going through the entire list is going to take me some time, probably weeks.

Maybe it is not worth-it enabling this until #438 is solved, but there are lots of cases this does improve, including quite a few cases in the closed source repos I work on in my day job, which I can't migrate to pyrefly until this is resolved.