Optional key to itertools.groupby

[achampion]

Redirect from: python/mypy#6377

itertools.groupby has an @overload type signature but this seems makes Optional[Callable[[T], S]] not match, e.g.:

def fn(iterable: Iterable[T], key: Optional[Callable[[T], S]] = None):
    iterator: Iterator[Tuple[S, Iterator[T]]] = itertools.groupby(iterable, key)
    ...

Fails with:

error: Argument 2 to "groupby" has incompatible type "Optional[Callable[[T], S]]"; expected "Callable[[T], S]"

The revealed type of itertools.groupby:

Overload(def [_T] (iterable: typing.Iterable[_T`-1]) -> typing.Iterator[Tuple[_T`-1, typing.Iterator[_T`-1]]], def [_T, _S] (iterable: typing.Iterable[_T`-1], key: def (_T`-1) -> _S`-2) -> typing.Iterator[Tuple[_S`-2, typing.Iterator[_T`-1]]])

It is legitimate to explicitly pass None as key, so is this expected behaviour?
You can do a similar @overload for myfunc to overcome the error.

If you use if key is not None: the problem becomes you can't effectively describe the type of iterator, e.g.:

iterator: Iterator[Tuple[S, Iterator[T]]]
if key is not None:
    iterator = itertools.groupby(iterable, key) # Ok
else:
    iterator = itertools.groupby(iterable)  # Fail

You get the error:

error: Incompatible types in assignment (expression has type "Iterator[Tuple[T, Iterator[T]]]", variable has type "Iterator[Tuple[S, Iterator[T]]]")

And you can't declare the types in each branch because you will get a redeclaration error Name 'iterator' already defined on line ...

Versions:
Python 3.7.2
MyPy 0.670

[srittau]

It seems that the seconds overload to groupby() is missing an Optional annotation. Patches welcome!