mypy: how to make it correctly guess the type?

user_5644961 0
Viewed 66

consider the following example

from typing import Callable, Iterable, TypeVar, Any, overload, Tuple, TYPE_CHECKING
from itertools import islice

T = TypeVar("T")

@overload
def mytake(n:int, iterable:Iterable[T]) -> Tuple[T,...]:...
@overload
def mytake(n:int, iterable:Iterable[T], container:Callable[[Iterable[T]],Any]) -> Any:...

def mytake(n:int, iterable:Iterable[T], container:Callable[[Iterable[T]],Any]=tuple) -> Any:
    return container(islice(iterable,n))

if TYPE_CHECKING:
    pass
else:
    reveal_type = print

#a:tuple[int,...]
a = mytake(10,range(100))
reveal_type(a)

#b:list[int]
b = mytake(10,range(100), list)
reveal_type(b)#desired to be list[int]

#c:int
c = mytake(10,range(100),sum)
reveal_type(c)#desired to be int


help(mytake)

that ouput

(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
45
Help on function mytake in module __main__:

mytake(n: int, iterable: Iterable[~T], container: Callable[[Iterable[~T]], Any] = <class 'tuple'>) -> Any

Basically just as desire, but I would like that mypy was the same

C:\Users\copperfield\Desktop\mypy_tests>mypy test6.py
test6.py:21: note: Revealed type is "builtins.tuple[builtins.int, ...]"
test6.py:25: note: Revealed type is "Any"
test6.py:29: note: Revealed type is "Any"
Success: no issues found in 1 source file

C:\Users\copperfield\Desktop\mypy_tests>

Beside type hinting the variables themselves, is there a way to make it so mypy correctly detect the types here?

python-3.10 mypy type-hinting python
1 Answers

Yes, with provisos. Overload resolution is pushing on the edge of mypy's already strained capabilities, so we can get close-ish.

What you want is an additional generic, for the return type of the callable.

T = TypeVar("T")
S = TypeVar("S")

@overload
def mytake(n:int, iterable: Iterable[T]) -> Tuple[T,...]:...
@overload
def mytake(n:int, iterable: Iterable[T], container: Callable[[Iterable[T]],S]) -> S:...

def mytake(n:int, iterable: Iterable[T], container: Optional[Callable[[Iterable[T]],S]] = None) -> S:
    if container is None:
        # We're in overload 1
        return cast(S, tuple(islice(iterable, n)))
    else:
        # We're in overload 2
        return container(islice(iterable, n))

Note that we have to cast in Case 1, as mypy doesn't consider the overloads that are available when type-checking the inside of the function. So we, as human readers, can conclude that, inside that if statement, we're in Overload 1, but mypy is still considering that container could be any Optional[Callable[[Iterable[T]], S]]. It's a safe cast, just one mypy doesn't quite get.

This is the signature we want, but it's awkward to work with.

#a:tuple[int,...]
a = mytake(10,range(100))
reveal_type(a)

The first one still works as intended.

#b:list[int]
b = mytake(10,range(100), list)
reveal_type(b)#desired to be list[int]

The second one won't work. We actually get an error. Mypy has some issues with interpreting types as Callable instances in generic contexts, so we can get around this using a lambda.

#b:list[int]
b = mytake(10,range(100), lambda x: list(x))
reveal_type(b)#desired to be list[int]

Not ideal, but workable.

The third one is not so lucky.

#c:int
c = mytake(10,range(100),sum)
reveal_type(c)#desired to be int

To be quite honest, I'm not even sure what happens here. Mypy reports that the type of c is Union[_T-1, builtins.int]. So a union of intand an unsolved metavariable. Lovely. I suspect this has to do with the fact thatsum` has several generic overloads and is probably just confusing mypy. Again, we can get around it with an explicit lambda.

#c:int
c = mytake(10,range(100),lambda x: sum(x))
reveal_type(c)#desired to be int

So the answer is yes, if you're willing to help the type checker out by generously inserting lambda everywhere. But this sort of thing pushes mypy to its limit. I'm not sure if other type checkers would do better generic inference in this case, as I'm most familiar with mypy. If someone has tried this with another Python type checker, feel free to post your results as well.

user_22886590
Related