Fix missing type store for overloads

[cdce8p]

Add missing call to store inferred types if an overload match is found early. All other code paths already do that.

Some background on the issue this fixes

I recently saw an interesting pattern in aiohttp to type values in an dict[str, Any] by subclassing dict.

T = TypeVar("T")
U = TypeVar("U")

class Key(Generic[T]):
    ...

class CustomDict(dict[Key[Any] | str, Any]):
    @overload  # type: ignore[override]
    def get(self, __key: Key[T]) -> T | None:
        ...

    @overload
    def get(self, __key: Key[T], __default: U) -> T | U:
        ...

    @overload
    def get(self, __key: str) -> Any | None:
        ...

    @overload
    def get(self, __key: str, __default: Any) -> Any:
        ...

    def get(self, __key: Key[Any] | str, __default: Any = None) -> Any:
        """Forward to super implementation."""
        return super().get(__key, __default)

    # overloads for __getitem__, setdefault, pop
    # ...

    @overload  # type: ignore[override]
    def __setitem__(self, key: Key[T], value: T) -> None:
        ...

    @overload
    def __setitem__(self, key: str, value: Any) -> None:
        ...

    def __setitem__(self, key: Key[Any] | str, value: Any) -> None:
        """Forward to super implementation."""
        return super().__setitem__(key, value)

With the exception that these overloads aren't technically compatible with the supertype, they do the job.

d = CustomDict()
key = Key[int]()
other_key = "other"
assert_type(d.get(key), int | None)
assert_type(d.get("other"), Any | None)

The issue exists for the __setitem__ case. Without this PR the following would create an issue. Here var would be inferred as dict[Never, Never], even though it should be dict[Any, Any] which is the case for non-subclassed dicts.

def a2(d: CustomDict) -> None:
    if (var := d.get("arg")) is None:
        var = d["arg"] = {}
        reveal_type(var)

[github-actions]

According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅

[hauntsaninja]

Nice, thanks!