Bidirectional inference for named expressions

[carljm]

We don't do multi-inference for walrus expressions, which causes bugs like this:

from typing import TypedDict, Any

class Inner(TypedDict):
    x: int

class Outer(TypedDict):
    inner: Inner

d1: Outer | dict[str, Any] = {"inner": {"x": 1}}
d2: Outer | dict[str, Any] = {"inner": (inner := {"x": 1})}

reveal_type(d1)  # Outer
reveal_type(d2)  # dict[str, Any]

https://play.ty.dev/f0dee682-cfe9-4c79-b6b9-6f7b7215f023

[AlexWaygood]

d2: Outer | dict[str, Any] = {"inner": (inner := {"x": 1})}  # errors

I don't see any error emitted on that line in the playground link you posted?

[carljm]

It was erroring when I created it! 😆

astral-sh/ruff#24151 changed the symptom here, but the underlying bug remains. The new symptom is just that we infer d2 as dict[str, Any] when really it should be Outer (and d1 is inferred as Outer).

Updated playground: https://play.ty.dev/f0dee682-cfe9-4c79-b6b9-6f7b7215f023

(Updated the OP also.)

There are probably other manifestations of the underlying issue that would result in diagnostics where they shouldn't be.

[ibraheemdev]

We actually don't support bidirectional inference at all with named expressions:

x: list[Any] = (y := [1])
reveal_type(x) # list[int]

Adding support is non-trivial because named expressions create a Definition, and we don't have type context support for Definition inference. This may require a similar change to scopes, where to produce the correct type of y we need to infer the outer scope to obtain the correct type context, but this will probably lead to cycles?

[carljm]

We could also have a dedicated Salsa query for inferring types of walrus-expression definitions, which takes an optional type context. If that query were called "in isolation" without type context, we'd get the context-free inference, but I would think that in actual fact we'd always just call it with the right type context, and not end up storing any additional memoized results? (Except in the actual multi-inference cases where we want to)

[ibraheemdev]

We could also have a dedicated Salsa query for inferring types of walrus-expression definitions, which takes an optional type context.

I'm not sure how that's different from adding a TypeContext input to infer_definition_types. We need to handle cases like:

def f(x: list[int | str]): ...

if f(x := [1]):
    reveal_type(x)

It's not clear to me how we would maintain the type context of list[int | str] when loading the place x without also inferring the outer expression. And if the outer expression is not standalone, that may require inferring the entire scope (which creates a cycle).

[carljm]

I'm not sure how that's different from adding a TypeContext input to infer_definition_types

It avoids the extra implicit interning layer that happens whenever a Salsa query has multiple arguments, for all non-walrus definitions. Maybe in practice this wouldn't be an issue, not sure.

It's not clear to me how we would maintain the type context of list[int | str] when loading the place x without also inferring the outer expression.

Yeah, good point, that is a problem. Having the in-context expression have a different type from the subsequent type of the walrus-defined name might be doable, but would be weird.

I think we don't need to prioritize this right now unless it shows up in user reports. And if/when it does, we should check the behavior of other type checkers in these cases.

[ibraheemdev]

It avoids the extra implicit interning layer that happens whenever a Salsa query has multiple arguments, for all non-walrus definitions. Maybe in practice this wouldn't be an issue, not sure.

(We use a salsa supertype trick to avoid interning in these cases, so that shouldn't be a problem)