Use type context from `__setitem__` signature to inform the type inferred on the right-hand side of subscript assignments

[nwalters512]

Summary

Consider this code:

from typing import TypedDict


class Bar(TypedDict):
    baz: float

foo: dict[str, Bar] = { "a": { "baz": 1 } }

foo["b"] = { "baz": 2 }

I would expect the assignment on the last line to be considered correct/valid, and this is how Pyright (playground link) and mypy (playground link) treat it. However, ty reports an error:

Invalid subscript assignment with key of type `Literal["b"]` and value of type `dict[Unknown | str, Unknown | int]` on object of type `dict[str, Bar]` (invalid-assignment) [Ln 9, Col 1]

The initial assignment/initialization of foo demonstrates that ty can correctly infer that a plain dictionary ({ "baz": 1 }) is assignable to type Bar; it just seemingly can't do that for a later assignment to a dictionary's items.

Playground link: https://play.ty.dev/8b35242f-36a7-4cb3-9e15-8416f0c509de

Version

ty 0.0.5 (d37b7dbd9 2025-12-20)

[AlexWaygood]

Thanks! Definitely a bug, though it might be a little tricky to fix.

When we see foo["b"] = { "baz": 2 }, we need to use the type context of the expected type for values of the foo dictionary to inform the way we infer the dictionary literal on the right-hand side of the assignment statement. The reason why this might be tricky is that to do this in a generalized way, we ideally would not special-case dictionaries, but treat all objects with __setitem__ methods the same way. So in terms of our implementation, I could imagine that this might look like:

1. Synthesize a type variable T
2. Ask the constraint solver what the best solution for T would be if type(foo).__setitem__(foo, "b", <instance of T>) must succeed. In this instance, it would hopefully solve T as being Bar
3. Given this solution for T, use that solution as type context to inform type inference of the dictionary literal on the right-hand side of this assignment

cc. @ibraheemdev for bidirectional inference

[carljm]

I'm not sure there'd be a need for synthesized typevars here, any more than anywhere else in bidirectional checking. I think this is just one case of the general issue (which I recall came up recently, though I can't seem to find an issue for it) that we should use type context for implicit calls to dunders just like we do for any other call.

[AlexWaygood]

I'm not an expert on our bidirectional-inference machinery so it's very possible I'm missing something, but how do we know what the type context for inferring the right-hand side even needs to be without invoking the constraint solver here? We have to infer a type for the right-hand side of the assignment before we pass that type into the call to __setitem__. But in order to infer the right-hand side of the assignment, we need to know what the type context needs to be in order for the call to __setitem__ to succeed.

That's why I mentioned we might first want to call __setitem__ with a synthesised TypeVar (without emitting diagnostics) to establish the type context, then infer the right-hand side of the assignment using that context, then after that call __setitem__ again with the inferred type of the right-hand side (with diagnostics switched on).

The other issue you were thinking of was probably https://github.com/astral-sh/ty/issues/2043

[carljm]

Yes, #2043 is the one I was thinking of.

we need to know what the type context needs to be in order for the call to setitem to succeed.

Isn't that just "the annotated type accepted by __setitem__"? I think that's generally how our bidirectional inference works for any call. It doesn't require calling the function with a synthesized typevar -- we just look at the parameter type on the callable, because we match arguments to parameters before doing type inference on the arguments. (Of course there are complexities with overloads, which then requires multi-inference.)

[AlexWaygood]

Okay cool, I was indeed missing something obvious :-) that sounds like a much easier way of doing things, yup!