Mypy complains about the `all_equal` recipe from the itertools docs

[AlexWaygood]

The following function is given as a recipe in the itertools docs, indicating that it's an idiomatic usage of groupby:

from itertools import groupby

def all_equal(iterable):
    g = groupby(iterable)
    return next(g, True) and not next(g, False)

I think the correct way of adding type annotations to this function would be as follows, since it will work on arbitrary iterables:

from collections.abc import Iterable
from itertools import groupby

def all_equal(iterable: Iterable[object]) -> bool:
    g = groupby(iterable)
    return next(g, True) and not next(g, False)

Unfortunately, however, mypy complains about this function:

error: Argument 1 to "next" has incompatible type "groupby[object, object]"; expected "SupportsNext[bool]"  [arg-type]

(Mypy gives a similar error if I use Iterable[Any] instead of Iterable[object] for the argument annotation.)

Perhaps we should consider copy-and-pasting all the itertools recipes into our test_cases directory. They're all meant to be idiomatic uses of itertools, so if any of them fail to type check, there's probably a problem somewhere.

[AlexWaygood]

(I encountered this in https://github.com/PyCQA/flake8-pyi/pull/432)

[AlexWaygood]

Hmm, I initially assumed this was probably an issue with typeshed's stubs -- either the stub for itertools.groupby, or the stub for builtins.next. Looking closer, I'm now not so sure. Pyright seems okay with this code; this may just be a problem with mypy's ability to solve the type variables here.

[srittau]

Perhaps we should consider copy-and-pasting all the itertools recipes into our test_cases directory. They're all meant to be idiomatic uses of itertools, so if any of them fail to type check, there's probably a problem somewhere.

Sounds like a good idea.

To the problem at hand: This sounds like a mypy bug to me. This is how the next() overload that should be used is defined:

@overload
def next(__i: SupportsNext[_T], __default: _VT) -> _T | _VT: ...

__i is groupby[object, object], which is compatible with SupportsNext[tuple[object, object]], so _T should be inferred as tuple[object, object] and _VT as bool. I'm not sure why mypy infers _T to be bool. It also seems mypy only complains about the first next() call (next(g, True)), not about the second (not next(g, False)).

[JelleZijlstra]

Probably mypy uses type context and thinks that the first next() call must return a bool. In fact its return type is always truthy and therefore ignored. It would arguably be cleaner to write the sample like this:

def all_equal(iterable: Iterable[object]) -> bool:
    g = groupby(iterable)
    next(g, True)
    return not next(g, False)

[AlexWaygood]

In fact its return type is always truthy and therefore ignored.

hah, good point -- arguably the itertools docs are being a little too clever for their own good there...