groupby with generic key function fails

[PiDelport]

Bug Report

Passing a generic function as key to groupby fails to type-check.

To Reproduce

This is the motivating example I encountered:

K = TypeVar("K")
V = TypeVar("V")

def fst(kv: Tuple[K, V]) -> K:
    k, v = kv
    return k


pairs = [(len(s), s) for s in ["one", "two", "three"]]

grouped = groupby(pairs, key=fst)

Mypy returns the following error:

error: Argument "key" to "groupby" has incompatible type "Callable[[Tuple[K, V]], K]"; expected "Callable[[Tuple[int, str]], K]"

I could simplify it down to the following:

T = TypeVar("T")

def identity(o: T) -> T:
    return o

grouped = groupby([1, 2, 3], key=identity)

Error:

error: Argument "key" to "groupby" has incompatible type "Callable[[T], T]"; expected "Callable[[int], T]"

In both cases, it works if the key function is modified to be non-generic.

Your Environment

• Mypy version used: mypy 0.790
• Python version used: 3.9

Workaround

In both cases, wrapping the generic function application in a lambda + application seems make the code type-check correctly:

grouped = groupby(pairs, key=lambda pair: fst(pair))

grouped = groupby([1, 2, 3], key=lambda o: identity(o))

This might be helpful for a fix, if Mypy can implicitly do whatever the above ends up doing?

[PiDelport]

This looks related to #6697 (map with abs), and #1317 (generic decorators & functions)?

[PiDelport]

I edited the issue with a lambda workaround: this could point at a way to fix this?

[erictraut]

I suspect this is a bug in mypy's constraint solver. It works fine in pyright.

I'll note that in the above example, the use of the TypeVar V in fst isn't really legitimate because it appears as a lone instance of that TypeVar in the function signature. It should be replaced by Any or some other non-generic type. Even if I correct this problem, mypy still emits a false positive error.

[hauntsaninja]

Yup, this is maybe my least favourite mypy issue. Thanks for the report, closing as a duplicate of the issues you mentioned.

Lambda's are somewhat special cased in mypy's inference logic, so not sure there's an easy way to translate that to a fix. I think #5738 is the issue that has the most constructive discussion on this subject.

[PiDelport]

@hauntsaninja: Is it worth keeping this open as an important / motivating use case, as per #6697 (comment) ?

[hauntsaninja]

I think it's covered / since the other issues are mentioned Github will link. It just needs to be fixed :-(