python / mypy

Optional static typing for Python
https://www.mypy-lang.org/
Other
18.24k stars 2.79k forks source link

Improve generic arguments inference for overloads #15750

Open eltoder opened 1 year ago

eltoder commented 1 year ago

Feature

In the following example (https://mypy-play.net/?mypy=latest&python=3.11&gist=3119048d3a61ecdda21806b02a0bd08c)

from typing import Any, Mapping, TypeVar, overload

_KT = TypeVar("_KT")
_VT = TypeVar("_VT")
_T1 = TypeVar("_T1")
_T2 = TypeVar("_T2")

class D(Mapping[_KT, _VT]):
    @overload
    def foo(self, value: Mapping[_KT, _VT]) -> D[_KT, _VT]: ...
    @overload
    def foo(self, value: Mapping[_T1, _T2]) -> D[_KT | _T1, _VT | _T2]: ...

    def foo(self, value: Any) -> Any:
        return self

def test(d: D[str, int]) -> None:
    d2: dict[str, Any] = {}
    reveal_type(d.foo(d2))
    d3: dict[int, Any] = {}
    reveal_type(d.foo(d3))

the types are revealed as

main.py:19: note: Revealed type is "__main__.D[Any, Any]"
main.py:21: note: Revealed type is "__main__.D[Union[builtins.str, builtins.int], Union[builtins.int, Any]]"
  1. The first one can be improved to D[str, Any] since both overloads return D[str, ...].
  2. For the second one, is there any point in infering Union[int, Any] instead of just Any? This is probably a general question, but I couldn't find any issues about it. Is there any context in which Union[int, Any] provides any extra information?
erictraut commented 1 year ago

For your second question, the answer is "yes int | Any provides more information than Any alone, and this can be valuable". For example, if you attempt to invoke a method on this value and that method isn't valid for an int, mypy can tell you about the potential bug. If the int subtype of the union is dropped, this bug will go unreported.

d.foo(d3)[0].capitalize()
eltoder commented 1 year ago

@erictraut ah, good point. If operations on unions are required to be valid for every component of the union, adding alternatives to Any is actually making the type more precise.