How to type hint a currying (partial) `map` function?

For simplicity assume map takes only one iterable.

I can type-declare a regular map function like this. (1) see end note

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

def regular_map(func: Callable[[T], S], iterable: Iterable[T]) -> Iterable[S]: 
    return map(func, iterable)

a = regular_map(lambda x: x, [1, 2, 3])
reveal_type(a)  # pyright: Type of "a" is "Iterable[int]"

But it doesn’t work when I try to write a partial map. (2)

def partial_map(func: Callable[[T], S]):
    def inner(iterable: Iterable[T]) -> Iterable[S]:
        return map(func, iterable)

    return inner

a = partial_map(lambda x: x)([1, 2, 3])
reveal_type(a)  # pyright: Type of "a" is "Iterable[Unknown]"

Actually, it works if the argument to partial_map is not a lambda but a type-hinted regular function. (3)

def foo(x: int) -> int:
    return x

a = partial(foo)([1, 2, 3])
reveal_type(a)  # pyright: Type of "a" is "Iterable[int]"

I want to provide type annotations for my partial map function, and not having it for lambdas is too much of a handicap. I guess it makes sense, lambdas don’t have any type annotations for type checkers to infer from. But then why does regular_map work? I’m partly curious, and partly hope there’s a secret to make partial_map work too.

Is what I want possible?


(P.S.) I tried with pyright and mypy. They behave similarly in case (1) and (2) but differ in (3), where pyright works as shown but mypy doesn’t, it fails to infer the type similarly to case (2). I don’t see why (3) should fail so I guess this is a mypy bug? Please correct me if I’m mistaken.

In your partial_map() definition, you didn’t actually define a type for the return value, so it’ll just be Any. You’d want to do Callable[[Iterable[T]], Iterable[S]]. With unhinted functions Pyright does attempt to infer types, so it might have done something and guessed, but Mypy refuses to do so.

2 Likes

With the Signature proposed by @TeamSpen210 , mypy is in fact able to correctly infer the type: mypy Playground
(although it starts generating funky error messages if you change the lambda away from the identify function)

pyright OTOH still struggles, I guess that is a limitation of their type inference.