python / mypy

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

Add support for simple lambda type guard #9890

Open charlax opened 3 years ago

charlax commented 3 years ago

Feature

Consider the following code:

from typing import Optional, List

list_with_none: List[Optional[int]] = [None, 1, 2]

filtered = filter(lambda v: v is not None, list_with_none)
reveal_type(filtered)
# Revealed type is 'typing.Iterator[Union[builtins.int, None]]'

It would be nice if mypy was able to parse simple lambda functions like this one, and infer that filtered is typing.Iterator[builtins.int].

Note that in this example, we use a simple none check, but this could work with more complicated checks:

list_of_tuples: List[Tuple[Optional[int], Optional[int]]]
filter(lambda v: v[0] is not None and v[1] is not None, list_of_tuples)
# would become `Iterator[Tuple[int, int]]`

Pitch

This is a fairly common idiom in Python.

Reasons why mypy might now want to do this:

gvanrossum commented 3 years ago

Let’s wait for PEP 647 rather than inventing something new.

gvanrossum commented 3 years ago

(I think once we have PEP 647 we can add an overload to filter() to support this.)

hauntsaninja commented 3 years ago

I don't see how this can be done with an overload, unless mypy learns to infer TypeGuards from lambdas (which doesn't necessarily seem the most straightforward)

gvanrossum commented 3 years ago

Hm... We already have this in builtins.pyi:

@overload
def filter(__function: None, __iterable: Iterable[Optional[_T]]) -> Iterator[_T]: ...
@overload
def filter(__function: Callable[[_T], Any], __iterable: Iterable[_T]) -> Iterator[_T]: ...

If we add an overload like you proposed earlier in https://github.com/python/mypy/pull/9865

def filter(fn: Callable[[T1], TypeGuard[T2]], it: Iterable[T1]) -> Iterator[T2]: ...

couldn't that be made to work? Not with a lambda, for sure -- you'd have to write a typeguard, e.g.

def not_none(x: Optional[T]) -> TypeGuard[T]:
    return x is not None

(It works with a filter() function defined like that locally -- I haven't yet tried adding that line to typeshed.)

hauntsaninja commented 3 years ago

Yes, sorry, I was talking specifically about lambdas since OP seemed specifically concerned about those. The typeshed overload should work for TypeGuard annotated functions :-)

AlexWaygood commented 2 years ago

I've changed the priority to "low", since the easy workaround here is just to define the function using a def statement instead of using a lambda.

(Note that the relevant overload was added to typeshed in https://github.com/python/typeshed/pull/6726.)